From 112c5a59e9113b8d1a9a52ccde2c5501bf862d38 Mon Sep 17 00:00:00 2001 From: zhichao-aws Date: Wed, 7 Feb 2024 15:38:55 +0800 Subject: [PATCH] ut for malformed response field Signed-off-by: zhichao-aws --- .../ml/engine/tools/MLModelToolTests.java | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/MLModelToolTests.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/MLModelToolTests.java index 5ff595713b..7969cae7e4 100644 --- a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/MLModelToolTests.java +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/MLModelToolTests.java @@ -98,6 +98,28 @@ public void testMLModelsWithDefaultOutputParserAndCustomizedResponseField() thro assertEquals("action1", future.get()); } + @Test + public void testMLModelsWithDefaultOutputParserAndMalformedResponseField() throws ExecutionException, InterruptedException { + ModelTensor modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("response", "response 1", "action", "action1")).build(); + ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build(); + ModelTensorOutput mlModelTensorOutput = ModelTensorOutput.builder().mlModelOutputs(Arrays.asList(modelTensors)).build(); + doAnswer(invocation -> { + + ActionListener actionListener = invocation.getArgument(2); + + actionListener.onResponse(MLTaskResponse.builder().output(mlModelTensorOutput).build()); + return null; + }).when(client).execute(eq(MLPredictionTaskAction.INSTANCE), any(), any()); + + Tool tool = MLModelTool.Factory.getInstance().create(Map.of("model_id", "modelId", "response_field", "malformed field")); + final CompletableFuture future = new CompletableFuture<>(); + ActionListener listener = ActionListener.wrap(r -> { future.complete(r); }, e -> { future.completeExceptionally(e); }); + tool.run(null, listener); + + future.join(); + assertEquals(null, future.get()); + } + @Test public void testMLModelsWithCustomizedOutputParser() { ModelTensor modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("thought", "thought 1", "action", "action1")).build();