fix bug - not found agent index (opensearch-project#1866)

Signed-off-by: Jing Zhang <[email protected]>
zane-neo · Jan 12, 2024 · 433d81b · 433d81b
1 parent e3a8a59
commit 433d81b
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java
@@ -170,8 +170,9 @@ public void execute(Input input, ActionListener<Output> listener) {
                     listener.onFailure(e);
                 }), context::restore));
             }
+        } else {
+            listener.onFailure(new ResourceNotFoundException("Agent index not found"));
         }
-
     }
 
     /**

diff --git a/...gorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java b/...gorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java
@@ -138,6 +138,25 @@ public void setup() {
 
     }
 
+    @Test
+    public void test_NoAgentIndex() {
+        ModelTensor modelTensor = ModelTensor.builder().name("response").dataAsMap(ImmutableMap.of("test_key", "test_value")).build();
+        Mockito.doAnswer(invocation -> {
+            ActionListener<ModelTensor> listener = invocation.getArgument(2);
+            listener.onResponse(modelTensor);
+            return null;
+        }).when(mlAgentRunner).run(Mockito.any(), Mockito.any(), Mockito.any());
+        Mockito.doReturn(mlAgentRunner).when(mlAgentExecutor).getAgentRunner(Mockito.any());
+        Mockito.when(metadata.hasIndex(Mockito.anyString())).thenReturn(false);
+
+        mlAgentExecutor.execute(getAgentMLInput(), agentActionListener);
+
+        Mockito.verify(agentActionListener).onFailure(exceptionCaptor.capture());
+        Exception exception = exceptionCaptor.getValue();
+        Assert.assertTrue(exception instanceof ResourceNotFoundException);
+        Assert.assertEquals(exception.getMessage(), "Agent index not found");
+    }
+
     @Test(expected = IllegalArgumentException.class)
     public void test_NullInput_ThrowsException() {
         mlAgentExecutor.execute(null, agentActionListener);