From 531cdc91214fc0463c15c5ad53d81e660257aaa0 Mon Sep 17 00:00:00 2001
From: Andriy Redko <andriy.redko@aiven.io>
Date: Mon, 11 Mar 2024 16:01:16 -0400
Subject: [PATCH] Cleanup SearchAsyncActionTests tests

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
---
 .../action/search/SearchAsyncActionTests.java | 33 ++++++++++---------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/server/src/test/java/org/opensearch/action/search/SearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/SearchAsyncActionTests.java
index 84e4f7a40bf70..35e90ff662b19 100644
--- a/server/src/test/java/org/opensearch/action/search/SearchAsyncActionTests.java
+++ b/server/src/test/java/org/opensearch/action/search/SearchAsyncActionTests.java
@@ -137,6 +137,10 @@ public void testSkipSearchShards() throws InterruptedException {
         lookup.put(replicaNode.getId(), new MockConnection(replicaNode));
         Map<String, AliasFilter> aliasFilters = Collections.singletonMap("_na_", new AliasFilter(null, Strings.EMPTY_ARRAY));
         AtomicInteger numRequests = new AtomicInteger(0);
+        final SearchRequestOperationsListener searchRequestOperationsListener = new SearchRequestOperationsListener.CompositeListener(
+            List.of(assertingListener),
+            LogManager.getLogger()
+        );
         AbstractSearchAsyncAction<TestSearchPhaseResult> asyncAction = new AbstractSearchAsyncAction<TestSearchPhaseResult>(
             "test",
             logger,
@@ -158,10 +162,7 @@ public void testSkipSearchShards() throws InterruptedException {
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(
-                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
-                request
-            ),
+            new SearchRequestContext(searchRequestOperationsListener, request),
             NoopTracer.INSTANCE
         ) {
 
@@ -193,7 +194,7 @@ protected SearchPhase getNextPhase(SearchPhaseResults<TestSearchPhaseResult> res
                     @Override
                     public void run() {
                         assertTrue(searchPhaseDidRun.compareAndSet(false, true));
-                        assertingListener.onPhaseEnd(new MockSearchPhaseContext(1, request, this), null);
+                        searchRequestOperationsListener.onPhaseEnd(new MockSearchPhaseContext(1, request, this), null);
                     }
                 };
             }
@@ -261,6 +262,10 @@ public void testLimitConcurrentShardRequests() throws InterruptedException {
         Map<String, AliasFilter> aliasFilters = Collections.singletonMap("_na_", new AliasFilter(null, Strings.EMPTY_ARRAY));
         CountDownLatch awaitInitialRequests = new CountDownLatch(1);
         AtomicInteger numRequests = new AtomicInteger(0);
+        final SearchRequestOperationsListener searchRequestOperationsListener = new SearchRequestOperationsListener.CompositeListener(
+            List.of(assertingListener),
+            LogManager.getLogger()
+        );
         AbstractSearchAsyncAction<TestSearchPhaseResult> asyncAction = new AbstractSearchAsyncAction<TestSearchPhaseResult>(
             "test",
             logger,
@@ -282,10 +287,7 @@ public void testLimitConcurrentShardRequests() throws InterruptedException {
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(
-                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
-                request
-            ),
+            new SearchRequestContext(searchRequestOperationsListener, request),
             NoopTracer.INSTANCE
         ) {
 
@@ -324,7 +326,7 @@ protected SearchPhase getNextPhase(SearchPhaseResults<TestSearchPhaseResult> res
                 return new SearchPhase("test") {
                     @Override
                     public void run() {
-                        assertingListener.onPhaseEnd(new MockSearchPhaseContext(1, request, this), null);
+                        searchRequestOperationsListener.onPhaseEnd(new MockSearchPhaseContext(1, request, this), null);
                         assertTrue(searchPhaseDidRun.compareAndSet(false, true));
                     }
                 };
@@ -630,6 +632,10 @@ public void testAllowPartialResults() throws InterruptedException {
         Map<String, AliasFilter> aliasFilters = Collections.singletonMap("_na_", new AliasFilter(null, Strings.EMPTY_ARRAY));
         AtomicInteger numRequests = new AtomicInteger(0);
         AtomicInteger numFailReplicas = new AtomicInteger(0);
+        final SearchRequestOperationsListener searchRequestOperationsListener = new SearchRequestOperationsListener.CompositeListener(
+            List.of(assertingListener),
+            LogManager.getLogger()
+        );
         AbstractSearchAsyncAction<TestSearchPhaseResult> asyncAction = new AbstractSearchAsyncAction<TestSearchPhaseResult>(
             "test",
             logger,
@@ -651,10 +657,7 @@ public void testAllowPartialResults() throws InterruptedException {
             new ArraySearchPhaseResults<>(shardsIter.size()),
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,
-            new SearchRequestContext(
-                new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
-                request
-            ),
+            new SearchRequestContext(searchRequestOperationsListener, request),
             NoopTracer.INSTANCE
         ) {
             @Override
@@ -688,7 +691,7 @@ protected SearchPhase getNextPhase(SearchPhaseResults<TestSearchPhaseResult> res
                     @Override
                     public void run() {
                         assertTrue(searchPhaseDidRun.compareAndSet(false, true));
-                        assertingListener.onPhaseEnd(new MockSearchPhaseContext(1, request, this), null);
+                        searchRequestOperationsListener.onPhaseEnd(new MockSearchPhaseContext(1, request, this), null);
                     }
                 };
             }