Cleanup AbstractSearchAsyncActionTests tests

Signed-off-by: Andriy Redko <[email protected]>
dzane17 · Mar 11, 2024 · d1c819e · d1c819e
1 parent b8d2731
commit d1c819e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 27 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchDfsQueryThenFetchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/SearchDfsQueryThenFetchAsyncAction.java
@@ -55,7 +55,7 @@
  *
  * @opensearch.internal
  */
-class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<DfsSearchResult> {
+final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<DfsSearchResult> {
 
     private final SearchPhaseController searchPhaseController;
 

diff --git a/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java
@@ -89,7 +89,6 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.instanceOf;
-import static org.mockito.Mockito.mock;
 
 public class AbstractSearchAsyncActionTests extends OpenSearchTestCase {
 
@@ -646,8 +645,8 @@ public void testOnPhaseListenersWithQueryAndThenFetchType() throws InterruptedEx
         assertEquals(1, testListener.getPhaseCurrent(expandPhase.getSearchPhaseName()));
 
         action.executeNextPhase(expandPhase, fetchPhase);
+        action.onPhaseDone(); /* finish phase since we don't have reponse being sent */
 
-        action.sendSearchResponse(mock(InternalSearchResponse.class), mock(String.valueOf(QuerySearchResult.class)));
         assertThat(testListener.getPhaseMetric(expandPhase.getSearchPhaseName()), greaterThanOrEqualTo(delay));
         assertEquals(1, testListener.getPhaseTotal(expandPhase.getSearchPhaseName()));
         assertEquals(0, testListener.getPhaseCurrent(expandPhase.getSearchPhaseName()));
@@ -673,14 +672,15 @@ public void testOnPhaseListenersWithDfsType() throws InterruptedException {
 
         searchDfsQueryThenFetchAsyncAction.skipShard(searchShardIterator);
         searchDfsQueryThenFetchAsyncAction.executeNextPhase(searchDfsQueryThenFetchAsyncAction, fetchPhase);
+        searchDfsQueryThenFetchAsyncAction.onPhaseFailure(
+            fetchPhase,
+            "Something went wrong",
+            null
+        ); /* finalizing the fetch phase since we do adhoc phase lifecycle calls */
 
         assertThat(testListener.getPhaseMetric(searchDfsQueryThenFetchAsyncAction.getSearchPhaseName()), greaterThanOrEqualTo(delay));
         assertEquals(1, testListener.getPhaseTotal(searchDfsQueryThenFetchAsyncAction.getSearchPhaseName()));
         assertEquals(0, testListener.getPhaseCurrent(searchDfsQueryThenFetchAsyncAction.getSearchPhaseName()));
-        searchDfsQueryThenFetchAsyncAction.sendSearchResponse(
-            mock(InternalSearchResponse.class),
-            mock(String.valueOf(QuerySearchResult.class))
-        );
     }
 
     private SearchDfsQueryThenFetchAsyncAction createSearchDfsQueryThenFetchAsyncAction(
@@ -719,7 +719,7 @@ private SearchDfsQueryThenFetchAsyncAction createSearchDfsQueryThenFetchAsyncAct
             null,
             null,
             null,
-            null,
+            controller,
             executor,
             resultConsumer,
             searchRequest,
@@ -734,18 +734,7 @@ private SearchDfsQueryThenFetchAsyncAction createSearchDfsQueryThenFetchAsyncAct
                 searchRequest
             ),
             NoopTracer.INSTANCE
-        ) {
-            @Override
-            public void sendSearchResponse(InternalSearchResponse internalSearchResponse, AtomicArray<SearchPhaseResult> queryResults) {
-                new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger).onPhaseEnd(
-                    this,
-                    new SearchRequestContext(
-                        new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger),
-                        searchRequest
-                    )
-                );
-            }
-        };
+        );
     }
 
     private SearchQueryThenFetchAsyncAction createSearchQueryThenFetchAsyncAction(
@@ -807,13 +796,7 @@ ShardSearchFailure[] buildShardFailures() {
 
             @Override
             public void sendSearchResponse(InternalSearchResponse internalSearchResponse, AtomicArray<SearchPhaseResult> queryResults) {
-                new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger).onPhaseEnd(
-                    this,
-                    new SearchRequestContext(
-                        new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger),
-                        searchRequest
-                    )
-                );
+                start();
             }
         };
     }