Add unit test for TransportUpdateModelAction class

Signed-off-by: Sicheng Song <[email protected]>
opensearch-project · Sep 23, 2023 · ac48986 · ac48986
1 parent 032579b
commit ac48986
Show file tree

Hide file tree

Showing 2 changed files with 263 additions and 52 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/action/models/TransportUpdateModelAction.java b/plugin/src/main/java/org/opensearch/ml/action/models/TransportUpdateModelAction.java
@@ -223,25 +223,6 @@ private void updateRemoteOrTextEmbeddingModel(
         }
     }
 
-    private ActionListener<UpdateResponse> getUpdateResponseListener(
-        String modelId,
-        ActionListener<UpdateResponse> actionListener,
-        ThreadContext.StoredContext context
-    ) {
-        return ActionListener.runBefore(ActionListener.wrap(updateResponse -> {
-            if (updateResponse != null && updateResponse.getResult() != DocWriteResponse.Result.UPDATED) {
-                log.info("Model id:{} failed update", modelId);
-                actionListener.onResponse(updateResponse);
-                return;
-            }
-            log.info("Completed Update Model Request, model id:{} updated", modelId);
-            actionListener.onResponse(updateResponse);
-        }, exception -> {
-            log.error("Failed to update ML model: " + modelId, exception);
-            actionListener.onFailure(exception);
-        }), context::restore);
-    }
-
     private void updateModelWithOrWithoutRelinkModelGroup(
         String modelId,
         String relinkModelGroupId,
@@ -271,4 +252,23 @@ private void updateModelWithOrWithoutRelinkModelGroup(
             client.update(updateRequest, getUpdateResponseListener(modelId, actionListener, context));
         }
     }
+
+    private ActionListener<UpdateResponse> getUpdateResponseListener(
+        String modelId,
+        ActionListener<UpdateResponse> actionListener,
+        ThreadContext.StoredContext context
+    ) {
+        return ActionListener.runBefore(ActionListener.wrap(updateResponse -> {
+            if (updateResponse != null && updateResponse.getResult() != DocWriteResponse.Result.UPDATED) {
+                log.info("Model id:{} failed update", modelId);
+                actionListener.onResponse(updateResponse);
+                return;
+            }
+            log.info("Completed Update Model Request, model id:{} updated", modelId);
+            actionListener.onResponse(updateResponse);
+        }, exception -> {
+            log.error("Failed to update ML model: " + modelId, exception);
+            actionListener.onFailure(exception);
+        }), context::restore);
+    }
 }