unit test

Signed-off-by: kkewwei <[email protected]>
opensearch-project · Mar 10, 2024 · a7aaa00 · a7aaa00
1 parent dd46b25
commit a7aaa00
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
@@ -306,7 +306,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
@@ -151,6 +151,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
             listener,
             controlled,
             false,
+            false,
             expected,
             new SearchShardIterator(null, null, Collections.emptyList(), null)
         );
@@ -162,6 +163,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
         ActionListener<SearchResponse> listener,
         final boolean controlled,
         final boolean failExecutePhaseOnShard,
+        final boolean catchExceptionInExecutePhaseOnShard,
         final AtomicLong expected,
         final SearchShardIterator... shards
     ) {
@@ -221,10 +223,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);
                     }
                 }
             }
@@ -513,6 +519,7 @@ public void onFailure(Exception e) {
             },
             false,
             true,
+            false,
             new AtomicLong(),
             shards
         );
@@ -559,6 +566,7 @@ public void onFailure(Exception e) {
             },
             false,
             false,
+            false,
             new AtomicLong(),
             shards
         );
@@ -574,7 +582,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))
@@ -610,6 +618,7 @@ public void onFailure(Exception e) {
             },
             false,
             false,
+            catchExceptionInExecutePhaseOnShard,
             new AtomicLong(),
             shards
         );
@@ -625,6 +634,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);