Add more UTs

Signed-off-by: zane-neo <zaniu@amazon.com>
zane-neo · Oct 27, 2023 · beb2011 · beb2011
1 parent 347fbfc
commit beb2011
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 0 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
@@ -207,4 +207,27 @@ public void executePredict_TextDocsInferenceInput() throws IOException {
         Assert.assertEquals(1, modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(0).getDataAsMap().size());
         Assert.assertEquals("value", modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(0).getDataAsMap().get("key"));
     }
+
+    @Test
+    public void test_initialize() {
+        ConnectorAction predictAction = ConnectorAction.builder()
+            .actionType(ConnectorAction.ActionType.PREDICT)
+            .method("POST")
+            .url("http://test.com/mock")
+            .requestBody("{\"input\": \"${parameters.input}\"}")
+            .build();
+        Map<String, String> credential = ImmutableMap.of(ACCESS_KEY_FIELD, encryptor.encrypt("test_key"), SECRET_KEY_FIELD, encryptor.encrypt("test_secret_key"));
+        Map<String, String> parameters = ImmutableMap.of(REGION_FIELD, "us-west-2", SERVICE_NAME_FIELD, "sagemaker");
+        Connector connector = AwsConnector.awsConnectorBuilder().name("test connector").version("1").protocol("http").parameters(parameters).credential(credential).actions(Arrays.asList(predictAction)).build();
+        connector.decrypt((c) -> encryptor.decrypt(c));
+        AwsConnectorExecutor executor = spy(new AwsConnectorExecutor(connector, httpClient));
+        initializeExecutor(executor);
+    }
+
+    private void initializeExecutor(RemoteConnectorExecutor executor) {
+        executor.setConnectionTimeoutInMillis(1000);
+        executor.setReadTimeoutInMillis(1000);
+        executor.setMaxConnections(30);
+        executor.initialize();
+    }
 }
diff --git a/...c/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java b/...c/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java
@@ -182,4 +182,24 @@ public void executePredict_TextDocsInput() throws IOException {
         Assert.assertArrayEquals(new Number[] {-0.014555434, -0.002135904, 0.0035105038}, modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(0).getData());
         Assert.assertArrayEquals(new Number[] {-0.014555434, -0.002135904, 0.0035105038}, modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(1).getData());
     }
+
+    @Test
+    public void test_initialize() {
+        ConnectorAction predictAction = ConnectorAction.builder()
+            .actionType(ConnectorAction.ActionType.PREDICT)
+            .method("POST")
+            .url("http://test.com/mock")
+            .requestBody("{\"input\": ${parameters.input}}")
+            .build();
+        Connector connector = HttpConnector.builder().name("test connector").version("1").protocol("http").actions(Arrays.asList(predictAction)).build();
+        HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector, httpClient));
+        initializeExecutor(executor);
+    }
+
+    private void initializeExecutor(RemoteConnectorExecutor executor) {
+        executor.setConnectionTimeoutInMillis(1000);
+        executor.setReadTimeoutInMillis(1000);
+        executor.setMaxConnections(30);
+        executor.initialize();
+    }
 }