From 568c7750894f098c8c268d2e56fd24c69c3fa712 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 8 Oct 2025 18:26:20 +0000 Subject: [PATCH] Add logic to throw exception on workload group deletion with associated rules (#19502) Signed-off-by: Kaushal Kumar (cherry picked from commit 2bb5e33af45048ada47d5c7ac00628dd5a83c8bb) Signed-off-by: github-actions[bot] --- CHANGELOG.md | 1 + .../plugin/wlm/WlmAutoTaggingIT.java | 2 +- .../plugin/wlm/WorkloadManagementPlugin.java | 2 +- .../TransportDeleteWorkloadGroupAction.java | 76 +++- .../WorkloadGroupPersistenceService.java | 10 +- .../wlm/WorkloadManagementPluginTests.java | 4 +- ...ansportDeleteWorkloadGroupActionTests.java | 326 +++++++++++++++++- 7 files changed, 392 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c82ce5c6e864..f2756540b5ae6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -109,6 +109,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Handle negative search request nodes stats ([#19340](https://github.com/opensearch-project/OpenSearch/pull/19340)) - Remove unnecessary iteration per-shard in request cache cleanup ([#19263](https://github.com/opensearch-project/OpenSearch/pull/19263)) - Fix derived field rewrite to handle range queries ([#19496](https://github.com/opensearch-project/OpenSearch/pull/19496)) +- [WLM] add a check to stop workload group deletion having rules ([#19502](https://github.com/opensearch-project/OpenSearch/pull/19502)) ### Dependencies - Bump `com.gradleup.shadow:shadow-gradle-plugin` from 8.3.5 to 8.3.9 ([#19400](https://github.com/opensearch-project/OpenSearch/pull/19400)) diff --git a/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmAutoTaggingIT.java b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmAutoTaggingIT.java index e8fd6520b41b9..4f8dfa89027ee 100644 --- a/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmAutoTaggingIT.java +++ b/plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmAutoTaggingIT.java @@ -784,7 +784,7 @@ public Collection createComponents( wlmClusterSettingValuesProvider, featureType ); - return List.of(refreshMechanism); + return List.of(refreshMechanism, rulePersistenceService, featureType); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java index d5bfede75926c..a6c18c964ed44 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/WorkloadManagementPlugin.java @@ -171,7 +171,7 @@ public Collection createComponents( wlmClusterSettingValuesProvider, featureType ); - return List.of(refreshMechanism); + return List.of(refreshMechanism, featureType, rulePersistenceService); } @Override diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java index 2bfbadba4d51d..bcaf4e868e4ff 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupAction.java @@ -8,6 +8,7 @@ package org.opensearch.plugin.wlm.action; +import org.opensearch.ResourceNotFoundException; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; @@ -15,15 +16,26 @@ import org.opensearch.cluster.block.ClusterBlockException; import org.opensearch.cluster.block.ClusterBlockLevel; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +import org.opensearch.rule.RulePersistenceService; +import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.autotagging.FeatureType; +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; import java.io.IOException; +import java.util.Collection; +import java.util.Collections; +import java.util.List; /** * Transport action for delete WorkloadGroup @@ -35,6 +47,8 @@ public class TransportDeleteWorkloadGroupAction extends TransportClusterManagerN AcknowledgedResponse> { private final WorkloadGroupPersistenceService workloadGroupPersistenceService; + private final RulePersistenceService rulePersistenceService; + private final FeatureType featureType; /** * Constructor for TransportDeleteWorkloadGroupAction @@ -45,6 +59,8 @@ public class TransportDeleteWorkloadGroupAction extends TransportClusterManagerN * @param threadPool - a {@link ThreadPool} object * @param indexNameExpressionResolver - a {@link IndexNameExpressionResolver} object * @param workloadGroupPersistenceService - a {@link WorkloadGroupPersistenceService} object + * @param persistenceService - a {@link IndexStoredRulePersistenceService} instance + * @param featureType - workloadManagement feature type */ @Inject public TransportDeleteWorkloadGroupAction( @@ -53,7 +69,9 @@ public TransportDeleteWorkloadGroupAction( ActionFilters actionFilters, ThreadPool threadPool, IndexNameExpressionResolver indexNameExpressionResolver, - WorkloadGroupPersistenceService workloadGroupPersistenceService + WorkloadGroupPersistenceService workloadGroupPersistenceService, + IndexStoredRulePersistenceService persistenceService, + WorkloadGroupFeatureType featureType ) { super( DeleteWorkloadGroupAction.NAME, @@ -65,6 +83,8 @@ public TransportDeleteWorkloadGroupAction( indexNameExpressionResolver ); this.workloadGroupPersistenceService = workloadGroupPersistenceService; + this.rulePersistenceService = persistenceService; + this.featureType = featureType; } @Override @@ -73,12 +93,18 @@ protected void clusterManagerOperation( ClusterState state, ActionListener listener ) throws Exception { - workloadGroupPersistenceService.deleteInClusterStateMetadata(request, listener); + threadPool.executor(executor()).submit(() -> { + try { + checkNoAssociatedRulesExist(request, listener, state); + } catch (Exception e) { + listener.onFailure(e); + } + }); } @Override protected String executor() { - return ThreadPool.Names.SAME; + return ThreadPool.Names.GET; } @Override @@ -90,4 +116,48 @@ protected AcknowledgedResponse read(StreamInput in) throws IOException { protected ClusterBlockException checkBlock(DeleteWorkloadGroupRequest request, ClusterState state) { return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_WRITE); } + + private void checkNoAssociatedRulesExist( + DeleteWorkloadGroupRequest request, + ActionListener listener, + ClusterState state + ) { + Collection workloadGroups = WorkloadGroupPersistenceService.getFromClusterStateMetadata(request.getName(), state); + if (workloadGroups.isEmpty()) { + throw new ResourceNotFoundException("No WorkloadGroup exists with the provided name: " + request.getName()); + } + + WorkloadGroup workloadGroup = workloadGroups.iterator().next(); + rulePersistenceService.getRule( + new GetRuleRequest(null, Collections.emptyMap(), null, featureType), + new ActionListener() { + @Override + public void onResponse(GetRuleResponse getRuleResponse) { + List associatedRules = getRuleResponse.getRules() + .stream() + .filter(rule -> rule.getFeatureValue().equals(workloadGroup.get_id())) + .toList(); + + if (!associatedRules.isEmpty()) { + listener.onFailure( + new IllegalStateException( + workloadGroup.getName() + + " workload group has rules with ids: " + + associatedRules + + " ." + + "Please delete them first otherwise system will be an inconsistent state." + ) + ); + return; + } + workloadGroupPersistenceService.deleteInClusterStateMetadata(request, listener); + } + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + } + ); + } } diff --git a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java index f37e90509c0fb..ecf5c7f68fa70 100644 --- a/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java +++ b/plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/WorkloadGroupPersistenceService.java @@ -236,14 +236,18 @@ protected AcknowledgedResponse newResponse(boolean acknowledged) { */ ClusterState deleteWorkloadGroupInClusterState(final String name, final ClusterState currentClusterState) { final Metadata metadata = currentClusterState.metadata(); - final WorkloadGroup workloadGroupToRemove = metadata.workloadGroups() + final WorkloadGroup workloadGroupToRemove = getWorkloadGroup(name, metadata); + + return ClusterState.builder(currentClusterState).metadata(Metadata.builder(metadata).remove(workloadGroupToRemove).build()).build(); + } + + private static WorkloadGroup getWorkloadGroup(String name, Metadata metadata) { + return metadata.workloadGroups() .values() .stream() .filter(workloadGroup -> workloadGroup.getName().equals(name)) .findAny() .orElseThrow(() -> new ResourceNotFoundException("No WorkloadGroup exists with the provided name: " + name)); - - return ClusterState.builder(currentClusterState).metadata(Metadata.builder(metadata).remove(workloadGroupToRemove).build()).build(); } /** diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java index 75cb5f01a6d40..81dc6baec5547 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/WorkloadManagementPluginTests.java @@ -50,6 +50,7 @@ import java.util.function.Supplier; import static org.opensearch.plugin.wlm.WorkloadManagementPlugin.PRINCIPAL_ATTRIBUTE_NAME; +import static org.hamcrest.Matchers.equalTo; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; @@ -187,8 +188,7 @@ public void testCreateComponentsReturnsRefreshMechanism() { mockRepositoriesServiceSupplier ); - assertEquals(1, components.size()); - assertTrue(components.iterator().next() instanceof RefreshBasedSyncMechanism); + assertThat(components.stream().filter(c -> c instanceof RefreshBasedSyncMechanism).count(), equalTo(1L)); } public void testSetAttributesWithMock() { diff --git a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java index 7ffa33aa8a80a..cffd6fffebea5 100644 --- a/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java +++ b/plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/action/TransportDeleteWorkloadGroupActionTests.java @@ -8,56 +8,344 @@ package org.opensearch.plugin.wlm.action; +import org.opensearch.ResourceNotFoundException; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.clustermanager.AcknowledgedResponse; import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.block.ClusterBlocks; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.WorkloadGroup; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.core.action.ActionListener; +import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType; import org.opensearch.plugin.wlm.service.WorkloadGroupPersistenceService; +import org.opensearch.rule.action.GetRuleRequest; +import org.opensearch.rule.action.GetRuleResponse; +import org.opensearch.rule.autotagging.Rule; +import org.opensearch.rule.service.IndexStoredRulePersistenceService; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; +import org.opensearch.wlm.MutableWorkloadGroupFragment; +import org.opensearch.wlm.ResourceType; +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; + +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; public class TransportDeleteWorkloadGroupActionTests extends OpenSearchTestCase { - ClusterService clusterService = mock(ClusterService.class); - TransportService transportService = mock(TransportService.class); - ActionFilters actionFilters = mock(ActionFilters.class); - ThreadPool threadPool = mock(ThreadPool.class); - IndexNameExpressionResolver indexNameExpressionResolver = mock(IndexNameExpressionResolver.class); - WorkloadGroupPersistenceService workloadGroupPersistenceService = mock(WorkloadGroupPersistenceService.class); + private ClusterService clusterService; + private TransportService transportService; + private ActionFilters actionFilters; + private ThreadPool threadPool; + private IndexNameExpressionResolver indexNameExpressionResolver; + private WorkloadGroupPersistenceService workloadGroupPersistenceService; + private IndexStoredRulePersistenceService rulePersistenceService; + private WorkloadGroupFeatureType featureType; + private TransportDeleteWorkloadGroupAction action; + + @Override + public void setUp() throws Exception { + super.setUp(); + clusterService = mock(ClusterService.class); + transportService = mock(TransportService.class); + actionFilters = mock(ActionFilters.class); + threadPool = mock(ThreadPool.class); + indexNameExpressionResolver = mock(IndexNameExpressionResolver.class); + workloadGroupPersistenceService = mock(WorkloadGroupPersistenceService.class); + rulePersistenceService = mock(IndexStoredRulePersistenceService.class); + featureType = mock(WorkloadGroupFeatureType.class); - TransportDeleteWorkloadGroupAction action = new TransportDeleteWorkloadGroupAction( - clusterService, - transportService, - actionFilters, - threadPool, - indexNameExpressionResolver, - workloadGroupPersistenceService - ); + action = new TransportDeleteWorkloadGroupAction( + clusterService, + transportService, + actionFilters, + threadPool, + indexNameExpressionResolver, + workloadGroupPersistenceService, + rulePersistenceService, + featureType + ); + } /** * Test case to validate the construction for TransportDeleteWorkloadGroupAction */ public void testConstruction() { assertNotNull(action); - assertEquals(ThreadPool.Names.SAME, action.executor()); + assertEquals(ThreadPool.Names.GET, action.executor()); } /** - * Test case to validate the clusterManagerOperation function in TransportDeleteWorkloadGroupAction + * Test case to validate successful workload group deletion */ - public void testClusterManagerOperation() throws Exception { - DeleteWorkloadGroupRequest request = new DeleteWorkloadGroupRequest("testGroup"); + public void testClusterManagerOperationSuccess() throws Exception { + String workloadGroupName = "testGroup"; + String workloadGroupId = "test-id-123"; + DeleteWorkloadGroupRequest request = new DeleteWorkloadGroupRequest(workloadGroupName); + @SuppressWarnings("unchecked") ActionListener listener = mock(ActionListener.class); - ClusterState clusterState = mock(ClusterState.class); + + // Create mock workload group + WorkloadGroup mockWorkloadGroup = createMockWorkloadGroup(workloadGroupName, workloadGroupId); + + // Create mock cluster state with workload group + ClusterState clusterState = createMockClusterStateWithWorkloadGroup(mockWorkloadGroup); + + // Mock empty rules response + GetRuleResponse getRuleResponse = mock(GetRuleResponse.class); + when(getRuleResponse.getRules()).thenReturn(Collections.emptyList()); + + // Mock executor service + ExecutorService mockExecutor = mock(ExecutorService.class); + when(threadPool.executor(ThreadPool.Names.GET)).thenReturn(mockExecutor); + doAnswer(invocation -> { + ((Runnable) invocation.getArgument(0)).run(); + return null; + }).when(mockExecutor).submit(any(Runnable.class)); + + // Mock rule persistence service responses + doAnswer(invocation -> { + @SuppressWarnings("unchecked") + ActionListener ruleListener = invocation.getArgument(1); + ruleListener.onResponse(getRuleResponse); + return null; + }).when(rulePersistenceService).getRule(any(GetRuleRequest.class), any()); + action.clusterManagerOperation(request, clusterState, listener); + verify(workloadGroupPersistenceService).deleteInClusterStateMetadata(eq(request), eq(listener)); + verify(mockExecutor).submit(any(Runnable.class)); + } + + /** + * Test case to validate ResourceNotFoundException when workload group doesn't exist + */ + public void testClusterManagerOperationWorkloadGroupNotFound() throws Exception { + String workloadGroupName = "nonExistentGroup"; + DeleteWorkloadGroupRequest request = new DeleteWorkloadGroupRequest(workloadGroupName); + + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + + // Create empty cluster state + ClusterState clusterState = createEmptyClusterState(); + + // Mock executor service + ExecutorService mockExecutor = mock(ExecutorService.class); + when(threadPool.executor(ThreadPool.Names.GET)).thenReturn(mockExecutor); + doAnswer(invocation -> { + ((Runnable) invocation.getArgument(0)).run(); + return null; + }).when(mockExecutor).submit(any(Runnable.class)); + + action.clusterManagerOperation(request, clusterState, listener); + + verify(listener).onFailure(any(ResourceNotFoundException.class)); + verify(workloadGroupPersistenceService, never()).deleteInClusterStateMetadata(any(), any()); + } + + /** + * Test case to validate that deletion is prevented when rules exist for the workload group + */ + public void testRuleDeletionWithExistingRules() throws Exception { + String workloadGroupName = "testGroup"; + String workloadGroupId = "test-id-123"; + DeleteWorkloadGroupRequest request = new DeleteWorkloadGroupRequest(workloadGroupName); + + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + + WorkloadGroup mockWorkloadGroup = createMockWorkloadGroup(workloadGroupName, workloadGroupId); + ClusterState clusterState = createMockClusterStateWithWorkloadGroup(mockWorkloadGroup); + + // Mock rules that reference this workload group + Rule mockRule = createMockRule("rule-1", workloadGroupId); + GetRuleResponse getRuleResponse = mock(GetRuleResponse.class); + when(getRuleResponse.getRules()).thenReturn(List.of(mockRule)); + + // Mock executor to immediately execute the runnable + ExecutorService mockExecutor = mock(ExecutorService.class); + when(threadPool.executor(ThreadPool.Names.GET)).thenReturn(mockExecutor); + doAnswer(invocation -> { + Runnable runnable = invocation.getArgument(0); + runnable.run(); + return null; + }).when(mockExecutor).submit(any(Runnable.class)); + + // Mock rule persistence service responses + doAnswer(invocation -> { + @SuppressWarnings("unchecked") + ActionListener ruleListener = invocation.getArgument(1); + ruleListener.onResponse(getRuleResponse); + return null; + }).when(rulePersistenceService).getRule(any(GetRuleRequest.class), any()); + + action.clusterManagerOperation(request, clusterState, listener); + + verify(rulePersistenceService).getRule(any(GetRuleRequest.class), any()); + verify(listener).onFailure(any(IllegalStateException.class)); + verify(workloadGroupPersistenceService, never()).deleteInClusterStateMetadata(any(), any()); + } + + /** + * Test case to validate successful deletion when no rules exist for the workload group + */ + public void testSuccessfulDeletionWithNoRules() throws Exception { + String workloadGroupName = "testGroup"; + String workloadGroupId = "test-id-123"; + DeleteWorkloadGroupRequest request = new DeleteWorkloadGroupRequest(workloadGroupName); + + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + + WorkloadGroup mockWorkloadGroup = createMockWorkloadGroup(workloadGroupName, workloadGroupId); + ClusterState clusterState = createMockClusterStateWithWorkloadGroup(mockWorkloadGroup); + + // Mock empty rules response + GetRuleResponse getRuleResponse = mock(GetRuleResponse.class); + when(getRuleResponse.getRules()).thenReturn(Collections.emptyList()); + + // Mock executor to immediately execute the runnable + ExecutorService mockExecutor = mock(ExecutorService.class); + when(threadPool.executor(ThreadPool.Names.GET)).thenReturn(mockExecutor); + doAnswer(invocation -> { + Runnable runnable = invocation.getArgument(0); + runnable.run(); + return null; + }).when(mockExecutor).submit(any(Runnable.class)); + + // Mock rule persistence service responses + doAnswer(invocation -> { + @SuppressWarnings("unchecked") + ActionListener ruleListener = invocation.getArgument(1); + ruleListener.onResponse(getRuleResponse); + return null; + }).when(rulePersistenceService).getRule(any(GetRuleRequest.class), any()); + + action.clusterManagerOperation(request, clusterState, listener); + + verify(rulePersistenceService).getRule(any(GetRuleRequest.class), any()); + verify(workloadGroupPersistenceService).deleteInClusterStateMetadata(eq(request), eq(listener)); + } + + /** + * Test case to validate rule deletion handles exceptions gracefully + */ + public void testRuleDeletionWithException() throws Exception { + String workloadGroupName = "testGroup"; + String workloadGroupId = "test-id-123"; + DeleteWorkloadGroupRequest request = new DeleteWorkloadGroupRequest(workloadGroupName); + + @SuppressWarnings("unchecked") + ActionListener listener = mock(ActionListener.class); + + WorkloadGroup mockWorkloadGroup = createMockWorkloadGroup(workloadGroupName, workloadGroupId); + ClusterState clusterState = createMockClusterStateWithWorkloadGroup(mockWorkloadGroup); + + ExecutorService mockExecutor = mock(ExecutorService.class); + when(threadPool.executor(ThreadPool.Names.GET)).thenReturn(mockExecutor); + doAnswer(invocation -> { + Runnable runnable = invocation.getArgument(0); + runnable.run(); + return null; + }).when(mockExecutor).submit(any(Runnable.class)); + + // Mock rule persistence service to throw exception + doAnswer(invocation -> { + @SuppressWarnings("unchecked") + ActionListener ruleListener = invocation.getArgument(1); + ruleListener.onFailure(new RuntimeException("Rule service error")); + return null; + }).when(rulePersistenceService).getRule(any(GetRuleRequest.class), any()); + + // Should not throw exception, just log and continue + action.clusterManagerOperation(request, clusterState, listener); + + verify(rulePersistenceService).getRule(any(GetRuleRequest.class), any()); + verify(listener).onFailure(any(RuntimeException.class)); + } + + /** + * Test case to validate cluster block check + */ + public void testCheckBlock() { + DeleteWorkloadGroupRequest request = new DeleteWorkloadGroupRequest("testGroup"); + ClusterState clusterState = mock(ClusterState.class); + ClusterBlocks clusterBlocks = mock(ClusterBlocks.class); + + when(clusterState.blocks()).thenReturn(clusterBlocks); + + action.checkBlock(request, clusterState); + + verify(clusterBlocks).globalBlockedException(any()); + } + + /** + * Test case to validate stream input reading + */ + public void testRead() throws IOException { + AcknowledgedResponse originalResponse = new AcknowledgedResponse(true); + + BytesStreamOutput out = new BytesStreamOutput(); + originalResponse.writeTo(out); + + AcknowledgedResponse readResponse = action.read(out.bytes().streamInput()); + + assertEquals(originalResponse.isAcknowledged(), readResponse.isAcknowledged()); + } + + private WorkloadGroup createMockWorkloadGroup(String name, String id) { + Map resourceLimits = Map.of(ResourceType.CPU, 0.5); + MutableWorkloadGroupFragment fragment = new MutableWorkloadGroupFragment( + MutableWorkloadGroupFragment.ResiliencyMode.ENFORCED, + resourceLimits + ); + return new WorkloadGroup(name, id, fragment, System.currentTimeMillis()); + } + + private ClusterState createMockClusterStateWithWorkloadGroup(WorkloadGroup workloadGroup) { + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + Map workloadGroups = Map.of(workloadGroup.get_id(), workloadGroup); + + when(clusterState.getMetadata()).thenReturn(metadata); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.workloadGroups()).thenReturn(workloadGroups); + + return clusterState; + } + + private ClusterState createEmptyClusterState() { + ClusterState clusterState = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + Map workloadGroups = Collections.emptyMap(); + + when(clusterState.getMetadata()).thenReturn(metadata); + when(clusterState.metadata()).thenReturn(metadata); + when(metadata.workloadGroups()).thenReturn(workloadGroups); + + return clusterState; + } + + private Rule createMockRule(String ruleId, String workloadGroupId) { + Rule mockRule = mock(Rule.class); + when(mockRule.getId()).thenReturn(ruleId); + when(mockRule.getFeatureValue()).thenReturn(workloadGroupId); + return mockRule; } }