Relax assertions for the wrapping

Author: pawankartik-elastic

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

@@ -89,6 +89,14 @@ protected boolean reuseClusters() {
         return false;
     }
 
+    private Exception unwrapIfWrappedInRemoteComputeException(Exception e) {
+        if (e instanceof RemoteComputeException rce) {
+            return (Exception) rce.getCause();
+        } else {
+            return e;
+        }
+    }
+
     private void createRemoteIndex(int numDocs) throws Exception {
         XContentBuilder mapping = JsonXContent.contentBuilder().startObject();
         mapping.startObject("runtime");
@@ -163,8 +171,9 @@ public void testCancel() throws Exception {
         } finally {
             SimplePauseFieldPlugin.allowEmitting.countDown();
         }
-        RemoteComputeException rce = expectThrows(RemoteComputeException.class, requestFuture::actionGet);
-        assertThat(rce.getCause().getMessage(), containsString("proxy timeout"));
+        Exception error = expectThrows(Exception.class, requestFuture::actionGet);
+        error = unwrapIfWrappedInRemoteComputeException(error);
+        assertThat(error.getCause().getMessage(), containsString("proxy timeout"));
     }
 
     public void testSameRemoteClusters() throws Exception {
@@ -284,7 +293,8 @@ public void testCancelSkipUnavailable() throws Exception {
             SimplePauseFieldPlugin.allowEmitting.countDown();
         }
 
-        RemoteComputeException error = expectThrows(RemoteComputeException.class, requestFuture::actionGet);
+        Exception error = expectThrows(Exception.class, requestFuture::actionGet);
+        error = unwrapIfWrappedInRemoteComputeException(error);
         assertThat(error.getCause(), instanceOf(TaskCancelledException.class));
     }
 }

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

@@ -64,6 +64,14 @@ protected Map<String, Boolean> skipUnavailableForRemoteClusters() {
         return Map.of(REMOTE_CLUSTER_1, randomBoolean(), REMOTE_CLUSTER_2, randomBoolean());
     }
 
+    private Exception unwrapIfWrappedInRemoteComputeException(Exception e) {
+        if (e instanceof RemoteComputeException rce) {
+            return (Exception) rce.getCause();
+        } else {
+            return e;
+        }
+    }
+
     public void testSuccessfulPathways() throws Exception {
         Map<String, Object> testClusterInfo = setupTwoClusters();
         int localNumShards = (Integer) testClusterInfo.get("local.num_shards");
@@ -816,11 +824,11 @@ public void testRemoteFailureSkipUnavailableTrue() throws IOException {
         String remote1Index = (String) testClusterInfo.get("remote.index");
         String q = Strings.format("FROM %s,cluster-a:%s*", localIndex, remote1Index);
 
-        RemoteComputeException rce = expectThrows(RemoteComputeException.class, () -> runQuery(q, false));
-        Throwable t = rce.getCause();
+        Exception error = expectThrows(Exception.class, () -> runQuery(q, false));
+        error = unwrapIfWrappedInRemoteComputeException(error);
 
-        assertThat(t, instanceOf(IllegalStateException.class));
-        assertThat(t.getMessage(), containsString("Accessing failing field"));
+        assertThat(error, instanceOf(IllegalStateException.class));
+        assertThat(error.getMessage(), containsString("Accessing failing field"));
     }
 
     private static void assertClusterMetadataInResponse(EsqlQueryResponse resp, boolean responseExpectMeta) {

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

@@ -99,18 +99,26 @@ private void assertClusterSuccess(EsqlQueryResponse resp, String clusterAlias, i
         assertThat(clusterInfo.getFailures(), empty());
     }
 
+    private Exception unwrapIfWrappedInRemoteComputeException(Exception e) {
+        if (e instanceof RemoteComputeException rce) {
+            return (Exception) rce.getCause();
+        } else {
+            return e;
+        }
+    }
+
     public void testPartialResults() throws Exception {
         populateIndices();
         EsqlQueryRequest request = new EsqlQueryRequest();
         request.query("FROM ok*,fail*,*:ok*,*:fail* | KEEP id, fail_me | LIMIT 1000");
         request.includeCCSMetadata(randomBoolean());
         {
             request.allowPartialResults(false);
-            RemoteComputeException rce = expectThrows(RemoteComputeException.class, () -> runQuery(request).close());
-            Throwable t = rce.getCause();
+            Exception error = expectThrows(Exception.class, () -> runQuery(request).close());
+            error = unwrapIfWrappedInRemoteComputeException(error);
 
-            assertThat(t, instanceOf(IllegalStateException.class));
-            assertThat(t.getMessage(), containsString("Accessing failing field"));
+            assertThat(error, instanceOf(IllegalStateException.class));
+            assertThat(error.getMessage(), containsString("Accessing failing field"));
         }
         request.allowPartialResults(true);
         try (var resp = runQuery(request)) {
@@ -194,8 +202,8 @@ public void sendResponse(Exception exception) {
             request.includeCCSMetadata(randomBoolean());
             {
                 request.allowPartialResults(false);
-                RemoteComputeException rce = expectThrows(RemoteComputeException.class, () -> runQuery(request).close());
-                var error = (Exception) rce.getCause();
+                Exception error = expectThrows(Exception.class, () -> runQuery(request).close());
+                error = unwrapIfWrappedInRemoteComputeException(error);
                 var unwrapped = ExceptionsHelper.unwrap(error, simulatedFailure.getClass());
                 assertNotNull(unwrapped);
                 assertThat(unwrapped.getMessage(), equalTo(simulatedFailure.getMessage()));
@@ -242,8 +250,8 @@ public void testFailToStartRequestOnRemoteCluster() throws Exception {
             request.includeCCSMetadata(randomBoolean());
             {
                 request.allowPartialResults(false);
-                var rce = expectThrows(RemoteComputeException.class, () -> runQuery(request).close());
-                var error = (Exception) rce.getCause();
+                Exception error = expectThrows(Exception.class, () -> runQuery(request).close());
+                error = unwrapIfWrappedInRemoteComputeException(error);
                 EsqlTestUtils.assertEsqlFailure(error);
                 var unwrapped = ExceptionsHelper.unwrap(error, simulatedFailure.getClass());
                 assertNotNull(unwrapped);