unit test

Signed-off-by: kkewwei <[email protected]>
opensearch-project · Feb 29, 2024 · 1b5b023 · 1b5b023
1 parent 5028a17
commit 1b5b023
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
@@ -305,7 +305,6 @@ public void onFailure(Exception t) {
                                 } else {
                                     onShardFailure(shardIndex, shard, shardIt, t);
                                 }
-
                             } finally {
                                 executeNext(pendingExecutions, thread);
                             }

diff --git a/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java
@@ -127,6 +127,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
             listener,
             controlled,
             false,
+            false,
             expected,
             new SearchShardIterator(null, null, Collections.emptyList(), null)
         );
@@ -138,6 +139,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
         ActionListener<SearchResponse> listener,
         final boolean controlled,
         final boolean failExecutePhaseOnShard,
+        final boolean catchExceptionInExecutePhaseOnShard,
         final AtomicLong expected,
         final SearchShardIterator... shards
     ) {
@@ -194,10 +196,14 @@ protected void executePhaseOnShard(
                 if (failExecutePhaseOnShard) {
                     listener.onFailure(new ShardNotFoundException(shardIt.shardId()));
                 } else {
-                    try {
+                    if (catchExceptionInExecutePhaseOnShard) {
+                        try {
+                            listener.onResponse(new QuerySearchResult());
+                        } catch (Exception e) {
+                            listener.onFailure(e);
+                        }
+                    } else {
                         listener.onResponse(new QuerySearchResult());
-                    } catch (Exception e) {
-                        listener.onFailure(e);
                     }
                 }
             }
@@ -517,6 +523,7 @@ public void onFailure(Exception e) {
             },
             false,
             true,
+            false,
             new AtomicLong(),
             shards
         );
@@ -563,6 +570,7 @@ public void onFailure(Exception e) {
             },
             false,
             false,
+            false,
             new AtomicLong(),
             shards
         );
@@ -578,7 +586,7 @@ public void onFailure(Exception e) {
         assertThat(searchResponse.getSuccessfulShards(), equalTo(shards.length));
     }
 
-    public void testExecutePhaseOnShardFailure() throws InterruptedException {
+    private void innerTestExecutePhaseOnShardFailure(boolean catchExceptionInExecutePhaseOnShard) throws InterruptedException {
         final Index index = new Index("test", UUID.randomUUID().toString());
 
         final SearchShardIterator[] shards = IntStream.range(0, 2 + randomInt(3))
@@ -614,6 +622,7 @@ public void onFailure(Exception e) {
             },
             false,
             false,
+            catchExceptionInExecutePhaseOnShard,
             new AtomicLong(),
             shards
         );
@@ -629,6 +638,14 @@ public void onFailure(Exception e) {
         assertThat(searchResponse.getSuccessfulShards(), equalTo(shards.length));
     }
 
+    public void testExecutePhaseOnShardFailure() throws InterruptedException {
+        innerTestExecutePhaseOnShardFailure(false);
+    }
+
+    public void testExecutePhaseOnShardFailureAndThrowException()  throws InterruptedException {
+        innerTestExecutePhaseOnShardFailure(true);
+    }
+
     public void testOnPhaseListenersWithQueryAndThenFetchType() throws InterruptedException {
         ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
         SearchRequestStats testListener = new SearchRequestStats(clusterSettings);