bug fix: to exclude remote model from deployment check (#4114) (#4116)

Author: opensearch-trigger-bot[bot]

plugin/src/main/java/org/opensearch/ml/helper/MemoryEmbeddingHelper.java

@@ -177,8 +177,8 @@ public void validateEmbeddingModelState(String modelId, FunctionName modelType,
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<MLModel> wrappedListener = ActionListener.runBefore(ActionListener.wrap(model -> {
                 MLModelState modelState = model.getModelState();
-                if (model.getAlgorithm() == FunctionName.REMOTE
-                    || (modelState != MLModelState.DEPLOYED && modelState != MLModelState.PARTIALLY_DEPLOYED)) {
+                if (model.getAlgorithm() != FunctionName.REMOTE
+                    && (modelState != MLModelState.DEPLOYED && modelState != MLModelState.PARTIALLY_DEPLOYED)) {
                     listener
                         .onFailure(
                             new IllegalStateException(

plugin/src/test/java/org/opensearch/ml/helper/MemoryEmbeddingHelperTests.java

@@ -293,6 +293,21 @@ public void testValidateEmbeddingModelStateDeployed() {
         verify(booleanListener).onResponse(true);
     }
 
+    @Test
+    public void testValidateEmbeddingModelStatePartiallyDeployed() {
+        when(mlModel.getModelState()).thenReturn(MLModelState.PARTIALLY_DEPLOYED);
+        
+        doAnswer(invocation -> {
+            ActionListener<MLModel> listener = invocation.getArgument(1);
+            listener.onResponse(mlModel);
+            return null;
+        }).when(mlModelManager).getModel(eq("model-123"), any());
+
+        helper.validateEmbeddingModelState("model-123", FunctionName.TEXT_EMBEDDING, booleanListener);
+
+        verify(booleanListener).onResponse(true);
+    }
+
     @Test
     public void testValidateEmbeddingModelStateNotDeployed() {
         when(mlModel.getModelState()).thenReturn(MLModelState.REGISTERED);
@@ -311,8 +326,46 @@ public void testValidateEmbeddingModelStateNotDeployed() {
         assertTrue(exceptionCaptor.getValue().getMessage().contains("DEPLOYED"));
     }
 
+    @Test
+    public void testValidateEmbeddingModelStateDeploying() {
+        when(mlModel.getModelState()).thenReturn(MLModelState.DEPLOYING);
+        
+        doAnswer(invocation -> {
+            ActionListener<MLModel> listener = invocation.getArgument(1);
+            listener.onResponse(mlModel);
+            return null;
+        }).when(mlModelManager).getModel(eq("model-123"), any());
+
+        helper.validateEmbeddingModelState("model-123", FunctionName.TEXT_EMBEDDING, booleanListener);
+
+        ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
+        verify(booleanListener).onFailure(exceptionCaptor.capture());
+        assertTrue(exceptionCaptor.getValue() instanceof IllegalStateException);
+        assertTrue(exceptionCaptor.getValue().getMessage().contains("DEPLOYED"));
+    }
+
     @Test
     public void testValidateEmbeddingModelStateRemoteModel() {
+        // Set up a model that has REMOTE algorithm but any state (e.g., REGISTERED)
+        when(mlModel.getAlgorithm()).thenReturn(FunctionName.REMOTE);
+        when(mlModel.getModelState()).thenReturn(MLModelState.REGISTERED);
+        
+        doAnswer(invocation -> {
+            ActionListener<MLModel> listener = invocation.getArgument(1);
+            listener.onResponse(mlModel);
+            return null;
+        }).when(mlModelManager).getModel(eq("model-123"), any());
+
+        // Pass TEXT_EMBEDDING as modelType to avoid early return, but model itself is REMOTE
+        helper.validateEmbeddingModelState("model-123", FunctionName.TEXT_EMBEDDING, booleanListener);
+        
+        verify(booleanListener).onResponse(true);
+        verify(mlModelManager).getModel(eq("model-123"), any());
+    }
+
+    @Test
+    public void testValidateEmbeddingModelStateRemoteModelTypeEarlyReturn() {
+        // Test the early return path when modelType itself is REMOTE
         helper.validateEmbeddingModelState("remote-model", FunctionName.REMOTE, booleanListener);
         verify(booleanListener).onResponse(true);
         verify(mlModelManager, never()).getModel(any(), any());