Skip to content

Commit 8e123f7

Browse files
ruai0511ruai0511
committed
add coverage
Signed-off-by: Ruirui Zhang <[email protected]>
1 parent c4bc6f5 commit 8e123f7

File tree

5 files changed

+132
-17
lines changed

5 files changed

+132
-17
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ public TransportCreateQueryGroupAction(
5151

5252
@Override
5353
protected void doExecute(Task task, CreateQueryGroupRequest request, ActionListener<CreateQueryGroupResponse> listener) {
54-
threadPool.executor(ThreadPool.Names.GENERIC)
54+
threadPool.executor(ThreadPool.Names.SAME)
5555
.execute(() -> queryGroupPersistenceService.persistInClusterStateMetadata(request.getQueryGroup(), listener));
5656
}
5757
}

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
* This class defines the functions for QueryGroup persistence
3535
*/
3636
public class QueryGroupPersistenceService {
37-
private static final String SOURCE = "query-group-persistence-service";
37+
static final String SOURCE = "query-group-persistence-service";
3838
private static final String CREATE_QUERY_GROUP_THROTTLING_KEY = "create-query-group";
3939
private static final Logger logger = LogManager.getLogger(QueryGroupPersistenceService.class);
4040
/**
@@ -62,7 +62,7 @@ public class QueryGroupPersistenceService {
6262
);
6363
private final ClusterService clusterService;
6464
private volatile int maxQueryGroupCount;
65-
private final ThrottlingKey createQueryGroupThrottlingKey;
65+
final ThrottlingKey createQueryGroupThrottlingKey;
6666

6767
/**
6868
* Constructor for QueryGroupPersistenceService
@@ -191,4 +191,11 @@ private Map<ResourceType, Double> calculateTotalUsage(Map<String, QueryGroup> ex
191191
}
192192
return map;
193193
}
194+
195+
/**
196+
* maxQueryGroupCount getter
197+
*/
198+
public int getMaxQueryGroupCount() {
199+
return maxQueryGroupCount;
200+
}
194201
}

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

Lines changed: 98 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,16 @@
1010

1111
import org.opensearch.cluster.ClusterName;
1212
import org.opensearch.cluster.ClusterState;
13+
import org.opensearch.cluster.ClusterStateUpdateTask;
1314
import org.opensearch.cluster.metadata.Metadata;
1415
import org.opensearch.cluster.metadata.QueryGroup;
1516
import org.opensearch.cluster.service.ClusterService;
1617
import org.opensearch.common.collect.Tuple;
1718
import org.opensearch.common.settings.ClusterSettings;
1819
import org.opensearch.common.settings.Settings;
20+
import org.opensearch.core.action.ActionListener;
21+
import org.opensearch.plugin.wlm.QueryGroupTestUtils;
22+
import org.opensearch.plugin.wlm.action.CreateQueryGroupResponse;
1923
import org.opensearch.search.ResourceType;
2024
import org.opensearch.test.OpenSearchTestCase;
2125
import org.opensearch.threadpool.ThreadPool;
@@ -26,6 +30,8 @@
2630
import java.util.List;
2731
import java.util.Map;
2832

33+
import org.mockito.ArgumentCaptor;
34+
2935
import static org.opensearch.cluster.metadata.QueryGroup.builder;
3036
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MEMORY_STRING;
3137
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.MONITOR_STRING;
@@ -34,18 +40,26 @@
3440
import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_ONE;
3541
import static org.opensearch.plugin.wlm.QueryGroupTestUtils._ID_TWO;
3642
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.assertEqualQueryGroups;
43+
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettings;
3744
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.clusterSettingsSet;
3845
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.preparePersistenceServiceSetup;
3946
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupList;
4047
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupOne;
4148
import static org.opensearch.plugin.wlm.QueryGroupTestUtils.queryGroupTwo;
4249
import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.QUERY_GROUP_COUNT_SETTING_NAME;
50+
import static org.opensearch.plugin.wlm.service.QueryGroupPersistenceService.SOURCE;
51+
import static org.mockito.ArgumentMatchers.any;
52+
import static org.mockito.ArgumentMatchers.eq;
53+
import static org.mockito.Mockito.anyString;
54+
import static org.mockito.Mockito.doAnswer;
4355
import static org.mockito.Mockito.mock;
56+
import static org.mockito.Mockito.times;
57+
import static org.mockito.Mockito.verify;
4458

4559
public class QueryGroupPersistenceServiceTests extends OpenSearchTestCase {
4660

4761
/**
48-
* Test case to validate the creation logic of a single QueryGroup
62+
* Test case to validate the creation logic of a QueryGroup
4963
*/
5064
public void testCreateQueryGroup() {
5165
Tuple<QueryGroupPersistenceService, ClusterState> setup = preparePersistenceServiceSetup(new HashMap<>());
@@ -136,6 +150,9 @@ public void testCreateQueryGroupOverflowCount() {
136150
assertThrows(RuntimeException.class, () -> queryGroupPersistenceService1.saveQueryGroupInClusterState(toCreate, clusterState));
137151
}
138152

153+
/**
154+
* Tests the invalid value of {@code node.query_group.max_count}
155+
*/
139156
public void testInvalidMaxQueryGroupCount() {
140157
Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 2).build();
141158
ClusterSettings clusterSettings = new ClusterSettings(settings, clusterSettingsSet());
@@ -147,4 +164,84 @@ public void testInvalidMaxQueryGroupCount() {
147164
);
148165
assertThrows(IllegalArgumentException.class, () -> queryGroupPersistenceService.setMaxQueryGroupCount(-1));
149166
}
167+
168+
/**
169+
* Tests the valid value of {@code node.query_group.max_count}
170+
*/
171+
public void testValidMaxSandboxCountSetting() {
172+
Settings settings = Settings.builder().put(QUERY_GROUP_COUNT_SETTING_NAME, 100).build();
173+
ClusterService clusterService = new ClusterService(settings, clusterSettings(), mock(ThreadPool.class));
174+
QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
175+
clusterService,
176+
settings,
177+
clusterSettings()
178+
);
179+
queryGroupPersistenceService.setMaxQueryGroupCount(50);
180+
assertEquals(50, queryGroupPersistenceService.getMaxQueryGroupCount());
181+
}
182+
183+
/**
184+
* Tests PersistInClusterStateMetadata function
185+
*/
186+
public void testPersistInClusterStateMetadata() {
187+
ClusterService clusterService = mock(ClusterService.class);
188+
@SuppressWarnings("unchecked")
189+
ActionListener<CreateQueryGroupResponse> listener = mock(ActionListener.class);
190+
QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
191+
clusterService,
192+
QueryGroupTestUtils.settings(),
193+
clusterSettings()
194+
);
195+
queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
196+
verify(clusterService).submitStateUpdateTask(eq(SOURCE), any());
197+
}
198+
199+
/**
200+
* Tests PersistInClusterStateMetadata function with inner functions
201+
*/
202+
public void testPersistInClusterStateMetadataInner() {
203+
ClusterService clusterService = mock(ClusterService.class);
204+
@SuppressWarnings("unchecked")
205+
ActionListener<CreateQueryGroupResponse> listener = mock(ActionListener.class);
206+
QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
207+
clusterService,
208+
QueryGroupTestUtils.settings(),
209+
clusterSettings()
210+
);
211+
ArgumentCaptor<ClusterStateUpdateTask> captor = ArgumentCaptor.forClass(ClusterStateUpdateTask.class);
212+
queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
213+
verify(clusterService, times(1)).submitStateUpdateTask(eq(SOURCE), captor.capture());
214+
ClusterStateUpdateTask capturedTask = captor.getValue();
215+
assertEquals(queryGroupPersistenceService.createQueryGroupThrottlingKey, capturedTask.getClusterManagerThrottlingKey());
216+
217+
doAnswer(invocation -> {
218+
ClusterStateUpdateTask task = invocation.getArgument(1);
219+
task.clusterStateProcessed(SOURCE, mock(ClusterState.class), mock(ClusterState.class));
220+
return null;
221+
}).when(clusterService).submitStateUpdateTask(anyString(), any());
222+
queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
223+
verify(listener).onResponse(any(CreateQueryGroupResponse.class));
224+
}
225+
226+
/**
227+
* Tests PersistInClusterStateMetadata function with failure
228+
*/
229+
public void testPersistInClusterStateMetadataFailure() {
230+
ClusterService clusterService = mock(ClusterService.class);
231+
@SuppressWarnings("unchecked")
232+
ActionListener<CreateQueryGroupResponse> listener = mock(ActionListener.class);
233+
QueryGroupPersistenceService queryGroupPersistenceService = new QueryGroupPersistenceService(
234+
clusterService,
235+
QueryGroupTestUtils.settings(),
236+
clusterSettings()
237+
);
238+
doAnswer(invocation -> {
239+
ClusterStateUpdateTask task = invocation.getArgument(1);
240+
Exception exception = new RuntimeException("Test Exception");
241+
task.onFailure(SOURCE, exception);
242+
return null;
243+
}).when(clusterService).submitStateUpdateTask(anyString(), any());
244+
queryGroupPersistenceService.persistInClusterStateMetadata(queryGroupOne, listener);
245+
verify(listener).onFailure(any(RuntimeException.class));
246+
}
150247
}

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

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
"test create a QueryGroup API successfully":
1+
"test create QueryGroup API":
22
- skip:
33
version: " - 2.16.99"
44
reason: "QueryGroup WorkloadManagement feature was added in 2.17"
@@ -20,12 +20,6 @@
2020
- match: { resource_limits.cpu: 0.4 }
2121
- match: { resource_limits.memory: 0.2 }
2222

23-
---
24-
"test bad QueryGroup API arguments":
25-
- skip:
26-
version: " - 2.16.99"
27-
reason: "QueryGroup WorkloadManagement feature was added in 2.17"
28-
2923
- do:
3024
catch: /illegal_argument_exception/
3125
create_query_group_context:
@@ -78,12 +72,6 @@
7872
}
7973
}
8074

81-
---
82-
"test create another QueryGroup API successfully":
83-
- skip:
84-
version: " - 2.16.99"
85-
reason: "QueryGroup WorkloadManagement feature was added in 2.17"
86-
8775
- do:
8876
create_query_group_context:
8977
body:

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,29 @@ public void testIllegalQueryGroupMode() {
103103
);
104104
}
105105

106+
public void testQueryGroupInitiation() {
107+
QueryGroup queryGroup = new QueryGroup("analytics", randomMode(), Map.of(ResourceType.MEMORY, 0.4));
108+
assertNotNull(queryGroup.getName());
109+
assertNotNull(queryGroup.get_id());
110+
assertNotNull(queryGroup.getResourceLimits());
111+
assertFalse(queryGroup.getResourceLimits().isEmpty());
112+
assertEquals(1, queryGroup.getResourceLimits().size());
113+
assertTrue(allowedModes.contains(queryGroup.getResiliencyMode()));
114+
assertTrue(queryGroup.getUpdatedAtInMillis() != 0);
115+
}
116+
117+
public void testIllegalQueryGroupName() {
118+
assertThrows(
119+
NullPointerException.class,
120+
() -> new QueryGroup("a".repeat(51), "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis())
121+
);
122+
assertThrows(
123+
NullPointerException.class,
124+
() -> new QueryGroup("", "_id", null, Map.of(ResourceType.MEMORY, 0.4), Instant.now().getMillis())
125+
);
126+
127+
}
128+
106129
public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() {
107130
assertThrows(
108131
IllegalArgumentException.class,

0 commit comments

Comments
 (0)