Address PR Feedback

Signed-off-by: Ryan Bogan <[email protected]>
opensearch-project · Dec 7, 2023 · 6f1a064 · 6f1a064
1 parent 80574a2
commit 6f1a064
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/src/main/java/org/opensearch/knn/common/KNNConstants.java b/src/main/java/org/opensearch/knn/common/KNNConstants.java
@@ -45,7 +45,7 @@ public class KNNConstants {
     public static final String MODEL_TIMESTAMP = "timestamp";
     public static final String MODEL_DESCRIPTION = "description";
     public static final String MODEL_ERROR = "error";
-    public static final String MODEL_NODE_ASSIGNMENT = "node_assignment";
+    public static final String MODEL_NODE_ASSIGNMENT = "training_node_assignment";
     public static final String PARAM_SIZE = "size";
     public static final Integer SEARCH_MODEL_MIN_SIZE = 1;
     public static final Integer SEARCH_MODEL_MAX_SIZE = 1000;

diff --git a/src/main/java/org/opensearch/knn/index/IndexUtil.java b/src/main/java/org/opensearch/knn/index/IndexUtil.java
@@ -34,7 +34,7 @@
 
 public class IndexUtil {
 
-    public static final String MODEL_NODE_ASSIGNMENT_KEY = "training_node_assignment";
+    public static final String MODEL_NODE_ASSIGNMENT_KEY = KNNConstants.MODEL_NODE_ASSIGNMENT;
 
     private static final Version MINIMAL_SUPPORTED_VERSION_FOR_IGNORE_UNMAPPED = Version.V_2_11_0;
     private static final Version MINIMAL_SUPPORTED_VERSION_FOR_MODEL_NODE_ASSIGNMENT = Version.V_2_12_0;

diff --git a/src/test/java/org/opensearch/knn/plugin/transport/GetModelResponseTests.java b/src/test/java/org/opensearch/knn/plugin/transport/GetModelResponseTests.java
@@ -46,7 +46,7 @@ public void testXContent() throws IOException {
         Model model = new Model(getModelMetadata(ModelState.CREATED), testModelBlob, modelId);
         GetModelResponse getModelResponse = new GetModelResponse(model);
         String expectedResponseString =
-            "{\"model_id\":\"test-model\",\"model_blob\":\"aGVsbG8=\",\"state\":\"created\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"node_assignment\":\"\"}";
+            "{\"model_id\":\"test-model\",\"model_blob\":\"aGVsbG8=\",\"state\":\"created\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}";
         XContentBuilder xContentBuilder = XContentFactory.jsonBuilder();
         getModelResponse.toXContent(xContentBuilder, null);
         assertEquals(expectedResponseString, xContentBuilder.toString());
@@ -57,7 +57,7 @@ public void testXContentWithNoModelBlob() throws IOException {
         Model model = new Model(getModelMetadata(ModelState.FAILED), null, modelId);
         GetModelResponse getModelResponse = new GetModelResponse(model);
         String expectedResponseString =
-            "{\"model_id\":\"test-model\",\"model_blob\":\"\",\"state\":\"failed\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"node_assignment\":\"\"}";
+            "{\"model_id\":\"test-model\",\"model_blob\":\"\",\"state\":\"failed\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}";
         XContentBuilder xContentBuilder = XContentFactory.jsonBuilder();
         getModelResponse.toXContent(xContentBuilder, null);
         assertEquals(expectedResponseString, xContentBuilder.toString());