Added another test for name check

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Oct 27, 2023 · a2eb324 · a2eb324
1 parent 864d0d7
commit a2eb324
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 17 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java
@@ -111,25 +111,25 @@ public void onFailure(Exception e) {
             registerModelGroupFuture.completeExceptionally(
                 new FlowFrameworkException("Model group name is not provided", RestStatus.BAD_REQUEST)
             );
-        }
+        } else {
+            MLRegisterModelGroupInputBuilder builder = MLRegisterModelGroupInput.builder();
+            builder.name(modelGroupName);
+            if (description != null) {
+                builder.description(description);
+            }
+            if (backendRoles != null && backendRoles.size() > 0) {
+                builder.backendRoles(backendRoles);
+            }
+            if (modelAccessMode != null) {
+                builder.modelAccessMode(modelAccessMode);
+            }
+            if (isAddAllBackendRoles != null) {
+                builder.isAddAllBackendRoles(isAddAllBackendRoles);
+            }
+            MLRegisterModelGroupInput mlInput = builder.build();
 
-        MLRegisterModelGroupInputBuilder builder = MLRegisterModelGroupInput.builder();
-        builder.name(modelGroupName);
-        if (description != null) {
-            builder.description(description);
+            mlClient.registerModelGroup(mlInput, actionListener);
         }
-        if (backendRoles != null && backendRoles.size() > 0) {
-            builder.backendRoles(backendRoles);
-        }
-        if (modelAccessMode != null) {
-            builder.modelAccessMode(modelAccessMode);
-        }
-        if (isAddAllBackendRoles != null) {
-            builder.isAddAllBackendRoles(isAddAllBackendRoles);
-        }
-        MLRegisterModelGroupInput mlInput = builder.build();
-
-        mlClient.registerModelGroup(mlInput, actionListener);
 
         return registerModelGroupFuture;
     }

diff --git a/src/test/java/org/opensearch/flowframework/workflow/ModelGroupStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/ModelGroupStepTests.java
@@ -36,6 +36,7 @@
 
 public class ModelGroupStepTests extends OpenSearchTestCase {
     private WorkflowData inputData = WorkflowData.EMPTY;
+    private WorkflowData inputDataWithNoName = WorkflowData.EMPTY;
 
     @Mock
     MachineLearningNodeClient machineLearningNodeClient;
@@ -54,6 +55,7 @@ public void setUp() throws Exception {
                 Map.entry("add_all_backend_roles", false)
             )
         );
+
     }
 
     public void testRegisterModelGroup() throws ExecutionException, InterruptedException, IOException {
@@ -103,4 +105,15 @@ public void testRegisterModelGroupFailure() throws ExecutionException, Interrupt
 
     }
 
+    public void testRegisterModelGroupWithNoName() throws IOException {
+        ModelGroupStep modelGroupStep = new ModelGroupStep(machineLearningNodeClient);
+
+        CompletableFuture<WorkflowData> future = modelGroupStep.execute(List.of(inputDataWithNoName));
+
+        assertTrue(future.isCompletedExceptionally());
+        ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());
+        assertTrue(ex.getCause() instanceof FlowFrameworkException);
+        assertEquals("Model group name is not provided", ex.getCause().getMessage());
+    }
+
 }