Fix failure UTs

Signed-off-by: zane-neo <[email protected]>
zane-neo · Feb 21, 2024 · ce9d409 · ce9d409
1 parent 189daf6
commit ce9d409
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/...ms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutorTest.java b/...ms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutorTest.java
@@ -163,8 +163,8 @@ public void executePredict_RemoteInferenceInput_EmptyIpAddress() {
         executor.executePredict(MLInput.builder().algorithm(FunctionName.REMOTE).inputDataset(inputDataSet).build(), actionListener);
         ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
         Mockito.verify(actionListener, times(1)).onFailure(exceptionCaptor.capture());
-        assert exceptionCaptor.getValue() instanceof MLException;
-        assertEquals("Fail to execute predict in aws connector", exceptionCaptor.getValue().getMessage());
+        assert exceptionCaptor.getValue() instanceof IllegalArgumentException;
+        assertEquals("Remote inference host name has private ip address: ", exceptionCaptor.getValue().getMessage());
     }
 
     @Test

diff --git a/...lgorithms/src/test/java/org/opensearch/ml/engine/httpclient/MLHttpClientFactoryTests.java b/...lgorithms/src/test/java/org/opensearch/ml/engine/httpclient/MLHttpClientFactoryTests.java
@@ -8,6 +8,7 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 
+import java.net.MalformedURLException;
 import java.util.Arrays;
 import java.util.Map;
 
@@ -172,6 +173,26 @@ public void test_validateIp_rarePrivateIp_throwException() throws Exception {
                 .actions(Arrays.asList(predictAction))
                 .build();
             ConnectorUtils.buildSdkRequest(connector, Map.of(), "hello world", SdkHttpMethod.POST);
+        } catch (MalformedURLException e) {
+            assertNotNull(e);
+        }
+
+        try {
+            ConnectorAction predictAction = ConnectorAction
+                .builder()
+                .actionType(ConnectorAction.ActionType.PREDICT)
+                .method("POST")
+                .url("http://153.24.76.232/mock")
+                .requestBody("{\"input\": \"${parameters.input}\"}")
+                .build();
+            Connector connector = HttpConnector
+                .builder()
+                .name("test connector")
+                .version("1")
+                .protocol("http")
+                .actions(Arrays.asList(predictAction))
+                .build();
+            ConnectorUtils.buildSdkRequest(connector, Map.of(), "hello world", SdkHttpMethod.POST);
         } catch (IllegalArgumentException e) {
             assertNotNull(e);
         }
@@ -242,8 +263,7 @@ public void test_validateSchemaAndPort_portNotInRange_throwException() throws Ex
 
     @Test
     public void test_validateSchemaAndPort_portNotANumber_throwException() throws Exception {
-        expectedException.expect(IllegalArgumentException.class);
-        expectedException.expectMessage("Port is not a valid number: abc");
+        expectedException.expect(MalformedURLException.class);
         ConnectorAction predictAction = ConnectorAction
             .builder()
             .actionType(ConnectorAction.ActionType.PREDICT)