From 433d81b0eb3a62e518387e2cbbe6572f5387a255 Mon Sep 17 00:00:00 2001 From: Jing Zhang Date: Fri, 12 Jan 2024 14:30:10 -0800 Subject: [PATCH] fix bug - not found agent index (#1866) Signed-off-by: Jing Zhang --- .../algorithms/agent/MLAgentExecutor.java | 3 ++- .../algorithms/agent/MLAgentExecutorTest.java | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 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 index 29c86bffc8..d0129e87ab 100644 --- 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 listener) { listener.onFailure(e); }), context::restore)); } + } else { + listener.onFailure(new ResourceNotFoundException("Agent index not found")); } - } /** diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java index d2ecb58f7a..f2e692f986 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java +++ b/ml-algorithms/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 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);