diff --git a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java index 9df460554..7e3da7c92 100644 --- 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 index 4ac187021..5833f0f22 100644 --- 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 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()); + } + }