Add comment

Author: dnhatn

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionBreakerIT.java

@@ -95,26 +95,31 @@ private EsqlQueryResponse runWithBreaking(EsqlQueryRequest request) throws Circu
 
     @Override
     protected EsqlQueryResponse run(EsqlQueryRequest request) {
+        if (randomBoolean()) {
+            request.allowPartialResults(randomBoolean());
+        }
+        Exception failure = null;
         try {
-            if (randomBoolean()) {
-                request.allowPartialResults(randomBoolean());
-            }
-            var resp = runWithBreaking(request);
+            final EsqlQueryResponse resp = runWithBreaking(request);
             if (resp.isPartial() == false) {
                 return resp;
             }
             try (resp) {
                 assertTrue(request.allowPartialResults());
             }
         } catch (Exception e) {
-            try (EsqlQueryResponse resp = super.run(request)) {
-                assertThat(e, instanceOf(CircuitBreakingException.class));
-                assertThat(ExceptionsHelper.status(e), equalTo(RestStatus.TOO_MANY_REQUESTS));
-                resp.incRef();
-                return resp;
+            failure = e;
+        }
+        // Re-run if the previous query failed or returned partial results
+        // Only check the previous failure if the second query succeeded
+        try (EsqlQueryResponse resp = super.run(request)) {
+            if (failure != null) {
+                assertThat(failure, instanceOf(CircuitBreakingException.class));
+                assertThat(ExceptionsHelper.status(failure), equalTo(RestStatus.TOO_MANY_REQUESTS));
             }
+            resp.incRef();
+            return resp;
         }
-        return super.run(request);
     }
 
     /**