diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/AbstractRetrieverToolTests.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/AbstractRetrieverToolTests.java
index f885a4867c..7f5b94f3cd 100644
--- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/AbstractRetrieverToolTests.java
+++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/AbstractRetrieverToolTests.java
@@ -125,27 +125,23 @@ public void testRunAsyncWithEmptySearchResponse() {
     }
 
     @Test
-    @SneakyThrows
     public void testRunAsyncWithIllegalQueryThenThrowException() {
         Client client = mock(Client.class);
         mockedImpl.setClient(client);
 
-        assertThrows(
-            "[input] is null or empty, can not process it.",
+        Exception exception = assertThrows(
             IllegalArgumentException.class,
             () -> mockedImpl.run(Map.of(AbstractRetrieverTool.INPUT_FIELD, ""), null)
         );
+        assertEquals("[input] is null or empty, can not process it.", exception.getMessage());
 
-        assertThrows(
-            "[input] is null or empty, can not process it.",
+        exception = assertThrows(
             IllegalArgumentException.class,
             () -> mockedImpl.run(Map.of(AbstractRetrieverTool.INPUT_FIELD, "  "), null)
         );
+        assertEquals("[input] is null or empty, can not process it.", exception.getMessage());
 
-        assertThrows(
-            "[input] is null or empty, can not process it.",
-            IllegalArgumentException.class,
-            () -> mockedImpl.run(Map.of("test", "hello world"), null)
-        );
+        exception = assertThrows(IllegalArgumentException.class, () -> mockedImpl.run(Map.of("test", "hello world"), null));
+        assertEquals("[input] is null or empty, can not process it.", exception.getMessage());
     }
 }