Fix UT

Signed-off-by: Heng Qian <[email protected]>
opensearch-project · Nov 21, 2024 · f779fd9 · f779fd9
1 parent 585e6d5
commit f779fd9
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/src/main/java/org/opensearch/agent/tools/CreateAlertTool.java b/src/main/java/org/opensearch/agent/tools/CreateAlertTool.java
@@ -66,6 +66,8 @@ public class CreateAlertTool implements Tool {
     @Getter
     private final String modelId;
     @Getter
+    private final String modelType;
+    @Getter
     private final String toolPrompt;
 
     private static final String MODEL_ID = "model_id";
@@ -93,9 +95,9 @@ public static ModelType from(String value) {
     public CreateAlertTool(Client client, String modelId, String modelType, String prompt) {
         this.client = client;
         this.modelId = modelId;
-        modelType = String.valueOf(ModelType.from(modelType));
+        this.modelType = String.valueOf(ModelType.from(modelType));
         if (prompt.isEmpty()) {
-            if (!promptDict.containsKey(modelType)) {
+            if (!promptDict.containsKey(this.modelType)) {
                 throw new IllegalArgumentException(
                     LoggerMessageFormat
                         .format(
@@ -106,7 +108,7 @@ public CreateAlertTool(Client client, String modelId, String modelType, String p
                         )
                 );
             }
-            this.toolPrompt = promptDict.get(modelType);
+            this.toolPrompt = promptDict.get(this.modelType);
         } else {
             this.toolPrompt = prompt;
         }

diff --git a/src/test/java/org/opensearch/agent/tools/CreateAlertToolTests.java b/src/test/java/org/opensearch/agent/tools/CreateAlertToolTests.java
@@ -159,16 +159,18 @@ public void testTool_WithBlankModelId() {
 
     @Test
     public void testTool_WithNonSupportedModelType() {
-        Exception exception = assertThrows(
-            IllegalArgumentException.class,
-            () -> CreateAlertTool.Factory
-                .getInstance()
-                .create(ImmutableMap.of("model_id", "modelId", "model_type", "non_supported_modelType"))
-        );
-        assertEquals(
-            "Failed to find the right prompt for modelType: non_supported_modelType, this tool supports prompts for these models: [CLAUDE,OPENAI]",
-            exception.getMessage()
-        );
+        CreateAlertTool alertTool = CreateAlertTool.Factory
+            .getInstance()
+            .create(ImmutableMap.of("model_id", "modelId", "model_type", "non_supported_modelType"));
+        assertEquals("CLAUDE", alertTool.getModelType());
+    }
+
+    @Test
+    public void testTool_WithEmptyModelType() {
+        CreateAlertTool alertTool = CreateAlertTool.Factory
+            .getInstance()
+            .create(ImmutableMap.of("model_id", "modelId", "model_type", ""));
+        assertEquals("CLAUDE", alertTool.getModelType());
     }
 
     @Test