|
@@ -6,6 +6,7 @@
|
|
|
*/
|
|
|
package org.elasticsearch.xpack.eql.analysis;
|
|
|
|
|
|
+import org.elasticsearch.Version;
|
|
|
import org.elasticsearch.action.ActionListener;
|
|
|
import org.elasticsearch.action.fieldcaps.FieldCapabilities;
|
|
|
import org.elasticsearch.action.fieldcaps.FieldCapabilitiesResponse;
|
|
@@ -18,9 +19,13 @@ import org.elasticsearch.cluster.ClusterName;
|
|
|
import org.elasticsearch.cluster.node.DiscoveryNode;
|
|
|
import org.elasticsearch.cluster.service.ClusterService;
|
|
|
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
|
|
|
+import org.elasticsearch.common.settings.Settings;
|
|
|
import org.elasticsearch.tasks.TaskCancelledException;
|
|
|
import org.elasticsearch.tasks.TaskId;
|
|
|
import org.elasticsearch.test.ESTestCase;
|
|
|
+import org.elasticsearch.test.transport.MockTransportService;
|
|
|
+import org.elasticsearch.threadpool.TestThreadPool;
|
|
|
+import org.elasticsearch.threadpool.ThreadPool;
|
|
|
import org.elasticsearch.transport.TransportService;
|
|
|
import org.elasticsearch.xpack.eql.action.EqlSearchRequest;
|
|
|
import org.elasticsearch.xpack.eql.action.EqlSearchResponse;
|
|
@@ -29,6 +34,8 @@ import org.elasticsearch.xpack.eql.execution.PlanExecutor;
|
|
|
import org.elasticsearch.xpack.eql.plugin.TransportEqlSearchAction;
|
|
|
import org.elasticsearch.xpack.ql.index.IndexResolver;
|
|
|
import org.elasticsearch.xpack.ql.type.DefaultDataTypeRegistry;
|
|
|
+import org.junit.After;
|
|
|
+import org.junit.Before;
|
|
|
import org.mockito.ArgumentCaptor;
|
|
|
import org.mockito.stubbing.Answer;
|
|
|
|
|
@@ -36,6 +43,7 @@ import java.util.Collections;
|
|
|
import java.util.HashMap;
|
|
|
import java.util.Map;
|
|
|
import java.util.concurrent.CountDownLatch;
|
|
|
+import java.util.concurrent.TimeUnit;
|
|
|
import java.util.concurrent.atomic.AtomicBoolean;
|
|
|
|
|
|
import static java.util.Collections.emptyMap;
|
|
@@ -51,17 +59,32 @@ import static org.mockito.Mockito.when;
|
|
|
|
|
|
public class CancellationTests extends ESTestCase {
|
|
|
|
|
|
+ private ThreadPool threadPool;
|
|
|
+ private TransportService transportService;
|
|
|
+
|
|
|
+ @Before
|
|
|
+ public void mockTransportService() {
|
|
|
+ threadPool = new TestThreadPool(getClass().getName());
|
|
|
+ // The TransportService needs to be able to return a valid RemoteClusterServices object down the stream, required by the Verifier.
|
|
|
+ transportService = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool);
|
|
|
+ }
|
|
|
+
|
|
|
+ @After
|
|
|
+ public void cleanupTransportService() {
|
|
|
+ ThreadPool.terminate(threadPool, 10, TimeUnit.SECONDS);
|
|
|
+ }
|
|
|
+
|
|
|
public void testCancellationBeforeFieldCaps() throws InterruptedException {
|
|
|
Client client = mock(Client.class);
|
|
|
EqlSearchTask task = mock(EqlSearchTask.class);
|
|
|
when(task.isCancelled()).thenReturn(true);
|
|
|
ClusterService mockClusterService = mockClusterService();
|
|
|
|
|
|
- IndexResolver indexResolver = new IndexResolver(client, randomAlphaOfLength(10), DefaultDataTypeRegistry.INSTANCE);
|
|
|
+ IndexResolver indexResolver = indexResolver(client);
|
|
|
PlanExecutor planExecutor = new PlanExecutor(client, indexResolver, new NamedWriteableRegistry(Collections.emptyList()));
|
|
|
CountDownLatch countDownLatch = new CountDownLatch(1);
|
|
|
TransportEqlSearchAction.operation(planExecutor, task, new EqlSearchRequest().query("foo where blah"), "",
|
|
|
- mock(TransportService.class), mockClusterService, new ActionListener<>() {
|
|
|
+ transportService, mockClusterService, new ActionListener<>() {
|
|
|
@Override
|
|
|
public void onResponse(EqlSearchResponse eqlSearchResponse) {
|
|
|
fail("Shouldn't be here");
|
|
@@ -120,11 +143,11 @@ public class CancellationTests extends ESTestCase {
|
|
|
}).when(client).fieldCaps(any(), any());
|
|
|
|
|
|
|
|
|
- IndexResolver indexResolver = new IndexResolver(client, randomAlphaOfLength(10), DefaultDataTypeRegistry.INSTANCE);
|
|
|
+ IndexResolver indexResolver = indexResolver(client);
|
|
|
PlanExecutor planExecutor = new PlanExecutor(client, indexResolver, new NamedWriteableRegistry(Collections.emptyList()));
|
|
|
CountDownLatch countDownLatch = new CountDownLatch(1);
|
|
|
TransportEqlSearchAction.operation(planExecutor, task, new EqlSearchRequest().indices("endgame")
|
|
|
- .query("process where foo==3"), "", mock(TransportService.class), mockClusterService, new ActionListener<>() {
|
|
|
+ .query("process where foo==3"), "", transportService, mockClusterService, new ActionListener<>() {
|
|
|
@Override
|
|
|
public void onResponse(EqlSearchResponse eqlSearchResponse) {
|
|
|
fail("Shouldn't be here");
|
|
@@ -185,11 +208,11 @@ public class CancellationTests extends ESTestCase {
|
|
|
return null;
|
|
|
}).when(client).execute(any(), searchRequestCaptor.capture(), any());
|
|
|
|
|
|
- IndexResolver indexResolver = new IndexResolver(client, randomAlphaOfLength(10), DefaultDataTypeRegistry.INSTANCE);
|
|
|
+ IndexResolver indexResolver = indexResolver(client);
|
|
|
PlanExecutor planExecutor = new PlanExecutor(client, indexResolver, new NamedWriteableRegistry(Collections.emptyList()));
|
|
|
CountDownLatch countDownLatch = new CountDownLatch(1);
|
|
|
TransportEqlSearchAction.operation(planExecutor, task, new EqlSearchRequest().indices("endgame")
|
|
|
- .query("process where foo==3"), "", mock(TransportService.class), mockClusterService, new ActionListener<>() {
|
|
|
+ .query("process where foo==3"), "", transportService, mockClusterService, new ActionListener<>() {
|
|
|
@Override
|
|
|
public void onResponse(EqlSearchResponse eqlSearchResponse) {
|
|
|
fail("Shouldn't be here");
|
|
@@ -227,4 +250,8 @@ public class CancellationTests extends ESTestCase {
|
|
|
when(mockClusterService.getClusterName()).thenReturn(mockClusterName);
|
|
|
return mockClusterService;
|
|
|
}
|
|
|
+
|
|
|
+ private static IndexResolver indexResolver(Client client) {
|
|
|
+ return new IndexResolver(client, randomAlphaOfLength(10), DefaultDataTypeRegistry.INSTANCE);
|
|
|
+ }
|
|
|
}
|