fix unit tests for abstract retriever tool

Signed-off-by: yuye-aws <[email protected]>
opensearch-project · Dec 12, 2023 · 36ab02a · 36ab02a
1 parent 41b33b9
commit 36ab02a
Showing 1 changed file with 6 additions and 10 deletions.
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
@@ -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());
     }
 }