Add asserting listener class

Signed-off-by: David Zane <[email protected]>
opensearch-project · Feb 26, 2024 · d9f86b2 · d9f86b2
1 parent d4c9e72
commit d9f86b2
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 49 deletions.
diff --git a/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java
@@ -86,8 +86,6 @@
 import java.util.function.BiFunction;
 import java.util.stream.IntStream;
 
-import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.CoreMatchers.nullValue;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.instanceOf;
@@ -107,27 +105,7 @@ public void setUp() throws Exception {
         super.setUp();
         executor = Executors.newFixedThreadPool(1);
         threadPool = new TestThreadPool(getClass().getName());
-        assertingListener = new SearchRequestOperationsListener() {
-            private volatile SearchPhase phase;
-
-            @Override
-            protected void onPhaseStart(SearchPhaseContext context) {
-                assertThat(phase, is(nullValue()));
-                phase = context.getCurrentPhase();
-            }
-
-            @Override
-            protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
-                assertThat(phase, is(context.getCurrentPhase()));
-                phase = null;
-            }
-
-            @Override
-            protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
-                assertThat(phase, is(context.getCurrentPhase()));
-                phase = null;
-            }
-        };
+        assertingListener = new SearchRequestOperationsListenerAssertingListener();
     }
 
     @After

diff --git a/server/src/test/java/org/opensearch/action/search/CanMatchPreFilterSearchPhaseTests.java b/server/src/test/java/org/opensearch/action/search/CanMatchPreFilterSearchPhaseTests.java
@@ -95,25 +95,7 @@ public void setUp() throws Exception {
         super.setUp();
 
         phases = Collections.newSetFromMap(new IdentityHashMap<>());
-        assertingListener = new SearchRequestOperationsListener() {
-            @Override
-            protected void onPhaseStart(SearchPhaseContext context) {
-                assertThat(phases.contains(context.getCurrentPhase()), is(false));
-                phases.add(context.getCurrentPhase());
-            }
-
-            @Override
-            protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
-                assertThat(phases.contains(context.getCurrentPhase()), is(true));
-                phases.remove(context.getCurrentPhase());
-            }
-
-            @Override
-            protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
-                assertThat(phases.contains(context.getCurrentPhase()), is(true));
-                phases.remove(context.getCurrentPhase());
-            }
-        };
+        assertingListener = new SearchRequestOperationsListenerAssertingListener();
     }
 
     @After
@@ -381,7 +363,7 @@ public void sendCanMatch(
                 randomIntBetween(1, 32),
                 SearchResponse.Clusters.EMPTY,
                 new SearchRequestContext(
-                    new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                    new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
                     searchRequest
                 ),
                 NoopTracer.INSTANCE

diff --git a/server/src/test/java/org/opensearch/action/search/SearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/SearchAsyncActionTests.java
@@ -80,6 +80,13 @@
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 
 public class SearchAsyncActionTests extends OpenSearchTestCase {
+    private SearchRequestOperationsListener assertingListener;
+
+    public void setUp() throws Exception {
+        super.setUp();
+
+        assertingListener = new SearchRequestOperationsListenerAssertingListener();
+    }
 
     public void testSkipSearchShards() throws InterruptedException {
         SearchRequest request = new SearchRequest();
@@ -139,7 +146,10 @@ public void testSkipSearchShards() throws InterruptedException {
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), request),
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
+                request
+            ),
             NoopTracer.INSTANCE
         ) {
 
@@ -259,7 +269,10 @@ public void testLimitConcurrentShardRequests() throws InterruptedException {
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), request),
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
+                request
+            ),
             NoopTracer.INSTANCE
         ) {
 
@@ -378,7 +391,10 @@ public void sendFreeContext(Transport.Connection connection, ShardSearchContextI
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), request),
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
+                request
+            ),
             NoopTracer.INSTANCE
         ) {
             TestSearchResponse response = new TestSearchResponse();
@@ -502,7 +518,10 @@ public void sendFreeContext(Transport.Connection connection, ShardSearchContextI
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), request),
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
+                request
+            ),
             NoopTracer.INSTANCE
         ) {
             TestSearchResponse response = new TestSearchResponse();
@@ -617,7 +636,10 @@ public void testAllowPartialResults() throws InterruptedException {
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), request),
+            new SearchRequestContext(
+                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
+                request
+            ),
             NoopTracer.INSTANCE
         ) {
             @Override

diff --git a/server/src/test/java/org/opensearch/action/search/SearchQueryThenFetchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/SearchQueryThenFetchAsyncActionTests.java
@@ -78,6 +78,14 @@
 import static org.hamcrest.Matchers.instanceOf;
 
 public class SearchQueryThenFetchAsyncActionTests extends OpenSearchTestCase {
+    private SearchRequestOperationsListener assertingListener;
+
+    public void setUp() throws Exception {
+        super.setUp();
+
+        assertingListener = new SearchRequestOperationsListenerAssertingListener();
+    }
+
     public void testBottomFieldSort() throws Exception {
         testCase(false, false);
     }
@@ -219,7 +227,7 @@ public void sendExecuteQuery(
             task,
             SearchResponse.Clusters.EMPTY,
             new SearchRequestContext(
-                new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()),
+                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
                 searchRequest
             ),
             NoopTracer.INSTANCE

diff --git a/...t/java/org/opensearch/action/search/SearchRequestOperationsListenerAssertingListener.java b/...t/java/org/opensearch/action/search/SearchRequestOperationsListenerAssertingListener.java
@@ -0,0 +1,35 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.search;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+public class SearchRequestOperationsListenerAssertingListener extends SearchRequestOperationsListener {
+    private volatile SearchPhase phase;
+
+    @Override
+    protected void onPhaseStart(SearchPhaseContext context) {
+        assertThat(phase, is(nullValue()));
+        phase = context.getCurrentPhase();
+    }
+
+    @Override
+    protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        assertThat(phase, is(context.getCurrentPhase()));
+        phase = null;
+    }
+
+    @Override
+    protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
+        assertThat(phase, is(context.getCurrentPhase()));
+        phase = null;
+    }
+}