Add test coverage for error cases.

Signed-off-by: Austin Lee <[email protected]>

Author: austintlee

search-processors/src/test/java/org/opensearch/searchpipelines/questionanswering/generative/llm/DefaultLlmImplTests.java

@@ -30,6 +30,8 @@
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
+import org.junit.Rule;
+import org.junit.rules.ExpectedException;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.opensearch.client.Client;
@@ -52,6 +54,9 @@ public class DefaultLlmImplTests extends OpenSearchTestCase {
     @Mock
     Client client;
 
+    @Rule
+    public ExpectedException exceptionRule = ExpectedException.none();
+
     public void testBuildMessageParameter() {
         DefaultLlmImpl connector = new DefaultLlmImpl("model_id", client);
         String question = "Who am I";
@@ -422,6 +427,47 @@ public void testChatCompletionBedrockThrowingError() throws Exception {
         assertEquals(errorMessage, (String) output.getErrors().get(0));
     }
 
+    public void testIllegalArgument1() {
+        exceptionRule.expect(IllegalArgumentException.class);
+        exceptionRule.expectMessage("Unknown/unsupported model provider: null.  You must provide a valid model provider or llm_response_field.");
+        MachineLearningInternalClient mlClient = mock(MachineLearningInternalClient.class);
+        ArgumentCaptor<MLInput> captor = ArgumentCaptor.forClass(MLInput.class);
+        DefaultLlmImpl connector = new DefaultLlmImpl("model_id", client);
+        connector.setMlClient(mlClient);
+
+        String errorMessage = "throttled";
+        Map<String, String> messageMap = Map.of("message", errorMessage);
+        Map<String, ?> dataAsMap = Map.of("error", messageMap);
+        ModelTensor tensor = new ModelTensor("tensor", new Number[0], new long[0], MLResultDataType.STRING, null, null, dataAsMap);
+        ModelTensorOutput mlOutput = new ModelTensorOutput(List.of(new ModelTensors(List.of(tensor))));
+        ActionFuture<MLOutput> future = mock(ActionFuture.class);
+        when(future.actionGet(anyLong())).thenReturn(mlOutput);
+        when(mlClient.predict(any(), any())).thenReturn(future);
+        ChatCompletionInput input = new ChatCompletionInput(
+            "model",
+            "question",
+            Collections.emptyList(),
+            Collections.emptyList(),
+            0,
+            "prompt",
+            "instructions",
+            null,
+            null
+        );
+        ChatCompletionOutput output = connector.doChatCompletion(input);
+    }
+
+    public void testIllegalArgument2() {
+        exceptionRule.expect(IllegalArgumentException.class);
+        exceptionRule.expectMessage("Unknown/unsupported model provider: null.  You must provide a valid model provider or llm_response_field.");
+        MachineLearningInternalClient mlClient = mock(MachineLearningInternalClient.class);
+        ArgumentCaptor<MLInput> captor = ArgumentCaptor.forClass(MLInput.class);
+        DefaultLlmImpl connector = new DefaultLlmImpl("model_id", client);
+        connector.setMlClient(mlClient);
+
+        connector.buildChatCompletionOutput(null, Collections.emptyMap(), null);
+    }
+
     private boolean isJson(String Json) {
         try {
             new JSONObject(Json);