add register connector

Signed-off-by: zhichao-aws <[email protected]>
zhichao-aws · Dec 29, 2023 · 90fb110 · 90fb110
1 parent d6fc397
commit 90fb110
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java b/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java
@@ -24,6 +24,7 @@
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.ml.common.MLModel;
 import org.opensearch.ml.common.MLTask;
 import org.opensearch.ml.common.MLTaskState;
 import org.opensearch.ml.common.input.execute.agent.AgentMLInput;
@@ -90,6 +91,12 @@ private Object parseFieldFromResponse(Response response, String field) {
         return result;
     }
 
+    protected String createConnector(String requestBody) {
+        Response response = makeRequest(client(), "POST", "/_plugins/_ml/connectors/_create", null, requestBody, null);
+        assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
+        return parseFieldFromResponse(response, MLModel.CONNECTOR_ID_FIELD).toString();
+    }
+
     protected String registerModel(String requestBody) {
         Response response = makeRequest(client(), "POST", "/_plugins/_ml/models/_register", null, requestBody, null);
         assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));