add coverage

Signed-off-by: Ruirui Zhang <mariazrr@amazon.com>

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/action/TransportCreateQueryGroupAction.java

@@ -44,14 +44,14 @@
         ThreadPool threadPool,
         QueryGroupPersistenceService queryGroupPersistenceService
     ) {
         super(CreateQueryGroupAction.NAME, transportService, actionFilters, CreateQueryGroupRequest::new);
         this.threadPool = threadPool;
         this.queryGroupPersistenceService = queryGroupPersistenceService;
     }
 
     @Override
     protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener<CreateQueryGroupResponse> listener) {
-        threadPool.executor(ThreadPool.Names.GENERIC)
+        threadPool.executor(ThreadPool.Names.SAME)
             .execute(() -> queryGroupPersistenceService.persistInClusterStateMetadata(request.getQueryGroup(), listener));
     }
 }

plugins/workload-management/src/main/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceService.java

@@ -34,7 +34,7 @@
  * This class defines the functions for QueryGroup persistence
  */
 public class QueryGroupPersistenceService {
-    private static final String SOURCE = "query-group-persistence-service";
+    static final String SOURCE = "query-group-persistence-service";
     private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group";
     private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class);
     /**
@@ -62,7 +62,7 @@
     );
     private final ClusterService clusterService;
     private volatile int maxQueryGroupCount;
-    private final ThrottlingKey createQueryGroupThrottlingKey;
+    final ThrottlingKey createQueryGroupThrottlingKey;
 
     /**
      * Constructor for QueryGroupPersistenceService
@@ -111,7 +111,7 @@
         clusterService.submitStateUpdateTask(SOURCE, new ClusterStateUpdateTask(Priority.NORMAL) {
             @Override
             public ClusterState execute(ClusterState currentState) throws Exception {
                 return saveQueryGroupInClusterState(queryGroup, currentState);
             }
 
             @Override
@@ -191,4 +191,11 @@
         }
         return map;
     }
+
+    /**
+     * maxQueryGroupCount getter
+     */
+    public int getMaxQueryGroupCount() {
+        return maxQueryGroupCount;
+    }
 }

plugins/workload-management/src/test/java/org/opensearch/plugin/wlm/service/QueryGroupPersistenceServiceTests.java

@@ -10,12 +10,16 @@
 
 import org.opensearch.cluster.ClusterName;
 import org.opensearch.cluster.ClusterState;
+import org.opensearch.cluster.ClusterStateUpdateTask;
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.metadata.QueryGroup;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.plugin.wlm.QueryGroupTestUtils;
+import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse;
 import org.opensearch.search.ResourceType;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -26,6 +30,8 @@
 import java.util.List;
 import java.util.Map;
 
+import org.mockito.ArgumentCaptor;
+
 import static org.opensearch.cluster.metadata.QueryGroup.builder;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MEMORY_STRING;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MONITOR_STRING;
@@ -34,18 +40,26 @@
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups;
+import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettings;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettingsSet;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne;
 import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo;
 import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.QUERY_GROUP_COUNT_SETTING_NAME;
+import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.SOURCE;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.anyString;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase {
 
     /**
-     * Test case to validate the creation logic of a single QueryGroup
+     * Test case to validate the creation logic of a QueryGroup
      */
     public void testCreateQueryGroup() {
         Tuple<QueryGroupPersistenceService, ClusterState> setup = preparePersistenceServiceSetup(new HashMap<>());
@@ -136,6 +150,9 @@ public void testCreateQueryGroupOverflowCount() {
         assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState));
     }
 
+    /**
+     * Tests the invalid value of {@code node.query_group.max_count}
+     */
     public void testInvalidMaxQueryGroupCount() {
         Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build();
         ClusterSettings clusterSettings = new ClusterSettings(settings, clusterSettingsSet());
@@ -147,4 +164,84 @@ public void testInvalidMaxQueryGroupCount() {
         );
         assertThrows(IllegalArgumentException.class, () -> queryGroupPersistenceService.setMaxQueryGroupCount(-1));
     }
+
+    /**
+     * Tests the valid value of {@code node.query_group.max_count}
+     */
+    public void testValidMaxSandboxCountSetting() {
+        Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 100).build();
+        ClusterService clusterService = new ClusterService(settings, clusterSettings(), mock(ThreadPool.class));
+        QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
+            clusterService,
+            settings,
+            clusterSettings()
+        );
+        queryGroupPersistenceService.setMaxQueryGroupCount(50);
+        assertEquals(50, queryGroupPersistenceService.getMaxQueryGroupCount());
+    }
+
+    /**
+     * Tests PersistInClusterStateMetadata function
+     */
+    public void testPersistInClusterStateMetadata() {
+        ClusterService clusterService = mock(ClusterService.class);
+        @SuppressWarnings("unchecked")
+        ActionListener<CreateQueryGroupResponse> listener = mock(ActionListener.class);
+        QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
+            clusterService,
+            QueryGroupTestUtils.settings(),
+            clusterSettings()
+        );
+        queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
+        verify(clusterService).submitStateUpdateTask(eq(SOURCE), any());
+    }
+
+    /**
+     * Tests PersistInClusterStateMetadata function with inner functions
+     */
+    public void testPersistInClusterStateMetadataInner() {
+        ClusterService clusterService = mock(ClusterService.class);
+        @SuppressWarnings("unchecked")
+        ActionListener<CreateQueryGroupResponse> listener = mock(ActionListener.class);
+        QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
+            clusterService,
+            QueryGroupTestUtils.settings(),
+            clusterSettings()
+        );
+        ArgumentCaptor<ClusterStateUpdateTask> captor = ArgumentCaptor.forClass(ClusterStateUpdateTask.class);
+        queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
+        verify(clusterService, times(1)).submitStateUpdateTask(eq(SOURCE), captor.capture());
+        ClusterStateUpdateTask capturedTask = captor.getValue();
+        assertEquals(queryGroupPersistenceService.createQueryGroupThrottlingKey, capturedTask.getClusterManagerThrottlingKey());
+
+        doAnswer(invocation -> {
+            ClusterStateUpdateTask task = invocation.getArgument(1);
+            task.clusterStateProcessed(SOURCE, mock(ClusterState.class), mock(ClusterState.class));
+            return null;
+        }).when(clusterService).submitStateUpdateTask(anyString(), any());
+        queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
+        verify(listener).onResponse(any(CreateQueryGroupResponse.class));
+    }
+
+    /**
+     * Tests PersistInClusterStateMetadata function with failure
+     */
+    public void testPersistInClusterStateMetadataFailure() {
+        ClusterService clusterService = mock(ClusterService.class);
+        @SuppressWarnings("unchecked")
+        ActionListener<CreateQueryGroupResponse> listener = mock(ActionListener.class);
+        QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
+            clusterService,
+            QueryGroupTestUtils.settings(),
+            clusterSettings()
+        );
+        doAnswer(invocation -> {
+            ClusterStateUpdateTask task = invocation.getArgument(1);
+            Exception exception = new RuntimeException("Test Exception");
+            task.onFailure(SOURCE, exception);
+            return null;
+        }).when(clusterService).submitStateUpdateTask(anyString(), any());
+        queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
+        verify(listener).onFailure(any(RuntimeException.class));
+    }
 }

plugins/workload-management/src/yamlRestTest/resources/rest-api-spec/test/wlm/10_create_query_group.yml

@@ -1,4 +1,4 @@
-"test create a QueryGroup API successfully":
+"test create QueryGroup API":
   - skip:
       version: " - 2.16.99"
       reason: "QueryGroup WorkloadManagement feature was added in 2.17"
@@ -20,12 +20,6 @@
   - match: { resource_limits.cpu: 0.4 }
   - match: { resource_limits.memory: 0.2 }
 
----
-"test bad QueryGroup API arguments":
-  - skip:
-      version: " - 2.16.99"
-      reason: "QueryGroup WorkloadManagement feature was added in 2.17"
-
   - do:
       catch: /illegal_argument_exception/
       create_query_group_context:
@@ -78,12 +72,6 @@
             }
           }
 
----
-"test create another QueryGroup API successfully":
-  - skip:
-      version: " - 2.16.99"
-      reason: "QueryGroup WorkloadManagement feature was added in 2.17"
-
   - do:
       create_query_group_context:
         body:

server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java

@@ -103,6 +103,29 @@ public void testIllegalQueryGroupMode() {
         );
     }
 
+    public void testQueryGroupInitiation() {
+        QueryGroup queryGroup = new QueryGroup("analytics", randomMode(), Map.of(ResourceType.MEMORY, 0.4));
+        assertNotNull(queryGroup.getName());
+        assertNotNull(queryGroup.get_id());
+        assertNotNull(queryGroup.getResourceLimits());
+        assertFalse(queryGroup.getResourceLimits().isEmpty());
+        assertEquals(1, queryGroup.getResourceLimits().size());
+        assertTrue(allowedModes.contains(queryGroup.getResiliencyMode()));
+        assertTrue(queryGroup.getUpdatedAtInMillis() != 0);
+    }
+
+    public void testIllegalQueryGroupName() {
+        assertThrows(
+            NullPointerException.class,
+            () -> new QueryGroup("a".repeat(51), "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis())
+        );
+        assertThrows(
+            NullPointerException.class,
+            () -> new QueryGroup("", "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis())
+        );
+
+    }
+
     public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() {
         assertThrows(
             IllegalArgumentException.class,