Adding model group id to integ tests (#189) (#190)

* Fix Model group id in integ tests Signed-off-by: Martin Gaievski <[email protected]> --------- Signed-off-by: Martin Gaievski <[email protected]> (cherry picked from commit 5228fc7) Co-authored-by: Martin Gaievski <[email protected]>
opensearch-project · Jun 1, 2023 · c75fc50 · c75fc50
1 parent 3e6aa00
commit c75fc50
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 0 deletions.
diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
@@ -9,6 +9,7 @@
 import static org.opensearch.neuralsearch.common.VectorUtil.vectorAsListToArray;
 
 import java.io.IOException;
+import java.net.URISyntaxException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Collections;
@@ -88,6 +89,9 @@ protected void updateClusterSettings(String settingKey, Object value) {
     }
 
     protected String uploadModel(String requestBody) throws Exception {
+        String modelGroupId = registerModelGroup();
+        // model group id is dynamically generated, we need to update model update request body after group is registered
+        requestBody = requestBody.replace("<MODEL_GROUP_ID>", modelGroupId);
         Response uploadResponse = makeRequest(
             client(),
             "POST",
@@ -492,4 +496,26 @@ protected static class KNNFieldConfig {
         private final Integer dimension;
         private final SpaceType spaceType;
     }
+
+    private String registerModelGroup() throws IOException, URISyntaxException {
+        String modelGroupRegisterRequestBody = Files.readString(
+            Path.of(classLoader.getResource("processor/CreateModelGroupRequestBody.json").toURI())
+        );
+        Response modelGroupResponse = makeRequest(
+            client(),
+            "POST",
+            "/_plugins/_ml/model_groups/_register",
+            null,
+            toHttpEntity(modelGroupRegisterRequestBody),
+            ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
+        );
+        Map<String, Object> modelGroupResJson = XContentHelper.convertToMap(
+            XContentFactory.xContent(XContentType.JSON),
+            EntityUtils.toString(modelGroupResponse.getEntity()),
+            false
+        );
+        String modelGroupId = modelGroupResJson.get("model_group_id").toString();
+        assertNotNull(modelGroupId);
+        return modelGroupId;
+    }
 }
diff --git a/src/test/resources/processor/CreateModelGroupRequestBody.json b/src/test/resources/processor/CreateModelGroupRequestBody.json
@@ -0,0 +1,4 @@
+{
+  "name": "test_model_group_public",
+  "description": "This is a public model group"
+}
diff --git a/src/test/resources/processor/UploadModelRequestBody.json b/src/test/resources/processor/UploadModelRequestBody.json
@@ -4,6 +4,7 @@
   "model_format": "TORCH_SCRIPT",
   "model_task_type": "text_embedding",
   "model_content_hash_value": "e13b74006290a9d0f58c1376f9629d4ebc05a0f9385f40db837452b167ae9021",
+  "model_group_id": "<MODEL_GROUP_ID>",
   "model_config": {
     "model_type": "bert",
     "embedding_dimension": 768,