add more unit tests

Signed-off-by: Chenyang Ji <cyji@amazon.com>
opensearch-project · Jul 9, 2024 · 720cefb · 720cefb
1 parent d00b2df
commit 720cefb
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/...s/query-insights/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java b/...s/query-insights/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java
@@ -12,6 +12,8 @@
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.util.Maps;
+import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo;
+import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.plugin.insights.rules.action.top_queries.TopQueries;
@@ -80,6 +82,25 @@ public static List<SearchQueryRecord> generateQueryInsightRecords(int lower, int
             attributes.put(Attribute.TOTAL_SHARDS, randomIntBetween(1, 100));
             attributes.put(Attribute.INDICES, randomArray(1, 3, Object[]::new, () -> randomAlphaOfLengthBetween(5, 10)));
             attributes.put(Attribute.PHASE_LATENCY_MAP, phaseLatencyMap);
+            attributes.put(
+                Attribute.TASK_RESOURCE_USAGES,
+                List.of(
+                    new TaskResourceInfo(
+                        randomAlphaOfLengthBetween(5, 10),
+                        randomLongBetween(1, 1000),
+                        randomLongBetween(1, 1000),
+                        randomAlphaOfLengthBetween(5, 10),
+                        new TaskResourceUsage(randomLongBetween(1, 1000), randomLongBetween(1, 1000))
+                    ),
+                    new TaskResourceInfo(
+                        randomAlphaOfLengthBetween(5, 10),
+                        randomLongBetween(1, 1000),
+                        randomLongBetween(1, 1000),
+                        randomAlphaOfLengthBetween(5, 10),
+                        new TaskResourceUsage(randomLongBetween(1, 1000), randomLongBetween(1, 1000))
+                    )
+                )
+            );
 
             records.add(new SearchQueryRecord(timestamp, measurements, attributes));
             timestamp += interval;

diff --git a/...rc/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java b/...rc/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java
@@ -32,6 +32,7 @@
 import org.opensearch.search.aggregations.support.ValueType;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.tasks.Task;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.test.ClusterServiceUtils;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.TestThreadPool;
@@ -65,6 +66,7 @@ public class QueryInsightsListenerTests extends OpenSearchTestCase {
     private final SearchRequest searchRequest = mock(SearchRequest.class);
     private final QueryInsightsService queryInsightsService = mock(QueryInsightsService.class);
     private final TopQueriesService topQueriesService = mock(TopQueriesService.class);
+    private final TaskResourceTrackingService taskResourceTrackingService = mock(TaskResourceTrackingService.class);
     private final ThreadPool threadPool = new TestThreadPool("QueryInsightsThreadPool");
     private ClusterService clusterService;
 
@@ -77,6 +79,7 @@ public void setup() {
         ClusterState state = ClusterStateCreationUtils.stateWithActivePrimary("test", true, 1 + randomInt(3), randomInt(2));
         clusterService = ClusterServiceUtils.createClusterService(threadPool, state.getNodes().getLocalNode(), clusterSettings);
         ClusterServiceUtils.setState(clusterService, state);
+        clusterService.setTaskResourceTrackingService(taskResourceTrackingService);
         when(queryInsightsService.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true);
         when(queryInsightsService.getTopQueriesService(MetricType.LATENCY)).thenReturn(topQueriesService);
 
@@ -139,6 +142,7 @@ public void testOnRequestEnd() throws InterruptedException {
         assertEquals(searchSourceBuilder.toString(), generatedRecord.getAttributes().get(Attribute.SOURCE));
         Map<String, String> labels = (Map<String, String>) generatedRecord.getAttributes().get(Attribute.LABELS);
         assertEquals("userLabel", labels.get(Task.X_OPAQUE_ID));
+        verify(taskResourceTrackingService, times(1)).refreshResourceStats(task);
     }
 
     public void testConcurrentOnRequestEnd() throws InterruptedException {
@@ -200,6 +204,7 @@ public void testConcurrentOnRequestEnd() throws InterruptedException {
         countDownLatch.await();
 
         verify(queryInsightsService, times(numRequests)).addRecord(any());
+        verify(taskResourceTrackingService, times(numRequests)).refreshResourceStats(task);
     }
 
     public void testSetEnabled() {