Refactor code and add additional ITs

Signed-off-by: Siddhant Deshmukh <[email protected]>
opensearch-project · Sep 4, 2024 · 9da0cbc · 9da0cbc
1 parent ff85d32
commit 9da0cbc
Show file tree

Hide file tree

Showing 5 changed files with 155 additions and 61 deletions.
diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsRestTestCase.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsRestTestCase.java
@@ -46,6 +46,7 @@
 import org.opensearch.core.xcontent.MediaType;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
 import org.opensearch.test.rest.OpenSearchRestTestCase;
 
 public abstract class QueryInsightsRestTestCase extends OpenSearchRestTestCase {
@@ -321,15 +322,57 @@ protected void waitForEmptyTopQueriesResponse() throws IOException, InterruptedE
         }
     }
 
-    protected String getTopQueries() throws IOException {
-        Request request = new Request("GET", "/_insights/top_queries?pretty");
+    protected void assertTopQueriesCount(int expectedTopQueriesCount, String type) throws IOException, InterruptedException {
+        // Ensure records are drained to the top queries service
+        Thread.sleep(QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL.millis());
+
+        // run five times to make sure the records are drained to the top queries services
+        for (int i = 0; i < 5; i++) {
+            String responseBody = getTopQueries(type);
+
+            int topNArraySize = countTopQueries(responseBody);
+
+            if (topNArraySize < expectedTopQueriesCount) {
+                // Ensure records are drained to the top queries service
+                Thread.sleep(QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL.millis());
+                continue;
+            }
+
+            // Validate that all queries are listed separately (no grouping)
+            Assert.assertEquals(expectedTopQueriesCount, topNArraySize);
+        }
+    }
+
+    protected String getTopQueries(String type) throws IOException {
+        // Base URL
+        String endpoint = "/_insights/top_queries?pretty";
+
+        if (type != null) {
+            switch (type.toLowerCase()) {
+                case "cpu":
+                case "memory":
+                case "latency":
+                    endpoint = "/_insights/top_queries?type=" + type + "&pretty";
+                    break;
+                case "all":
+                    // Keep the default endpoint (no type parameter)
+                    break;
+                default:
+                    // Throw an exception if the type is invalid
+                    throw new IllegalArgumentException("Invalid type: " + type + ". Valid types are 'all', 'cpu', 'memory', or 'latency'.");
+            }
+        }
+
+        Request request = new Request("GET", endpoint);
         Response response = client().performRequest(request);
+
         Assert.assertEquals(200, response.getStatusLine().getStatusCode());
 
         String responseBody = new String(response.getEntity().getContent().readAllBytes(), StandardCharsets.UTF_8);
         return responseBody;
     }
 
+
     protected void updateClusterSettings(Supplier<String> settingsSupplier) throws IOException {
         Request request = new Request("PUT", "/_cluster/settings");
         request.setJsonEntity(settingsSupplier.get());

diff --git a/.../java/org/opensearch/plugin/insights/core/service/grouper/MinMaxQueryGrouperByNoneIT.java b/.../java/org/opensearch/plugin/insights/core/service/grouper/MinMaxQueryGrouperByNoneIT.java
@@ -33,23 +33,7 @@ public void testGroupingByNone() throws IOException, InterruptedException {
         doSearch("match", 6);
         doSearch("term", 4);
 
-        // Ensure records are drained to the top queries service
-        Thread.sleep(QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL.millis());
-
-        // run five times to make sure the records are drained to the top queries services
-        for (int i = 0; i < 5; i++) {
-            String responseBody = getTopQueries();
-
-            int topNArraySize = countTopQueries(responseBody);
-
-            if (topNArraySize == 0) {
-                Thread.sleep(QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL.millis());
-                continue;
-            }
-
-            // Validate that all queries are listed separately (no grouping)
-            Assert.assertEquals(12, topNArraySize);
-        }
+        assertTopQueriesCount(12, "latency");
     }
 
     private String groupByNoneSettings() {

diff --git a/...org/opensearch/plugin/insights/core/service/grouper/MinMaxQueryGrouperBySimilarityIT.java b/...org/opensearch/plugin/insights/core/service/grouper/MinMaxQueryGrouperBySimilarityIT.java
@@ -36,20 +36,7 @@ public void testGroupingBySimilarity() throws IOException, InterruptedException
         doSearch("match", 6);
         doSearch("term", 4);
 
-        // run five times to make sure the records are drained to the top queries services
-        for (int i = 0; i < 5; i++) {
-            String responseBody = getTopQueries();
-
-            int topNArraySize = countTopQueries(responseBody);
-
-            if (topNArraySize == 0) {
-                Thread.sleep(QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL.millis());
-                continue;
-            }
-
-            // Validate that all queries are grouped
-            Assert.assertEquals(3, topNArraySize);
-        }
+        assertTopQueriesCount(3, "latency");
     }
 
     /**

diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxQueryGrouperIT.java b/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxQueryGrouperIT.java
@@ -0,0 +1,102 @@
+package org.opensearch.plugin.insights.core.service.grouper;
+
+import org.junit.Assert;
+import org.opensearch.plugin.insights.QueryInsightsRestTestCase;
+import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
+
+import java.io.IOException;
+
+public class MinMaxQueryGrouperIT extends QueryInsightsRestTestCase {
+    /**
+     * Grouping by none should not group queries
+     * @throws IOException
+     * @throws InterruptedException
+     */
+    public void testNoneToSimilarityGroupingTransition() throws IOException, InterruptedException {
+
+        waitForEmptyTopQueriesResponse();
+
+        updateClusterSettings(this::defaultTopQueriesSettings);
+
+        // Search
+        doSearch("range", 2);
+        doSearch("match", 6);
+        doSearch("term", 4);
+
+        assertTopQueriesCount(12, "latency");
+
+        updateClusterSettings(this::defaultTopQueryGroupingSettings);
+
+        // Top queries should be drained due to grouping change from NONE -> SIMILARITY
+        assertTopQueriesCount(0, "latency");
+
+        // Search
+        doSearch("range", 2);
+        doSearch("match", 6);
+        doSearch("term", 4);
+
+        // 3 groups
+        assertTopQueriesCount(3, "latency");
+    }
+
+    public void testSimilarityToNoneGroupingTransition() throws IOException, InterruptedException {
+
+        waitForEmptyTopQueriesResponse();
+
+        updateClusterSettings(this::defaultTopQueryGroupingSettings);
+
+        // Search
+        doSearch("range", 2);
+        doSearch("match", 6);
+        doSearch("term", 4);
+
+        assertTopQueriesCount(3, "latency");
+
+        updateClusterSettings(this::defaultTopQueriesSettings);
+
+        // Top queries should be drained due to grouping change from SIMILARITY -> NONE
+        assertTopQueriesCount(0, "latency");
+
+        // Search
+        doSearch("range", 2);
+        doSearch("match", 6);
+        doSearch("term", 4);
+
+        assertTopQueriesCount(12, "latency");
+    }
+
+    public void testSimilarityMaxGroupsChanged() throws IOException, InterruptedException {
+
+        waitForEmptyTopQueriesResponse();
+
+        updateClusterSettings(this::defaultTopQueryGroupingSettings);
+
+        // Search
+        doSearch("range", 2);
+        doSearch("match", 6);
+        doSearch("term", 4);
+
+        assertTopQueriesCount(3, "latency");
+
+        // Change max groups exluding topn setting
+        updateClusterSettings(this::updateMaxGroupsExcludingTopNSetting);
+
+        // Top queries should be drained due to max group change
+        assertTopQueriesCount(0, "latency");
+
+        // Search
+        doSearch("range", 2);
+        doSearch("match", 6);
+        doSearch("term", 4);
+
+        assertTopQueriesCount(3, "latency");
+    }
+
+    protected String updateMaxGroupsExcludingTopNSetting() {
+        return "{\n"
+            + "    \"persistent\" : {\n"
+            + "        \"search.insights.top_queries.max_groups_excluding_topn\" : 1\n"
+            + "    }\n"
+            + "}";
+    }
+}
diff --git a/...t/java/org/opensearch/plugin/insights/rules/resthandler/top_queries/TopQueriesRestIT.java b/...t/java/org/opensearch/plugin/insights/rules/resthandler/top_queries/TopQueriesRestIT.java
@@ -11,6 +11,8 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+
 import org.junit.Assert;
 import org.opensearch.client.Request;
 import org.opensearch.client.Response;
@@ -55,39 +57,15 @@ public void testTopQueriesResponses() throws IOException, InterruptedException {
 
         doSearch(2);
 
-        // run five times to make sure the records are drained to the top queries services
-        for (int i = 0; i < 5; i++) {
-            String responseBody = getTopQueries();
-
-            int topNArraySize = countTopQueries(responseBody);
-
-            if (topNArraySize == 0) {
-                Thread.sleep(QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL.millis());
-                continue;
-            }
-
-            Assert.assertEquals(2, topNArraySize);
-        }
+        assertTopQueriesCount(2, "latency");
 
         // Enable Top N Queries by resource usage
         updateClusterSettings(this::topQueriesByResourceUsagesSettings);
 
         // Do Search
         doSearch(2);
 
-        // Run five times to make sure the records are drained to the top queries services
-        for (int i = 0; i < 5; i++) {
-            String responseBody = getTopQueries();
-
-            int topNArraySize = countTopQueries(responseBody);
-
-            if (topNArraySize == 0) {
-                Thread.sleep(QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL.millis());
-                continue;
-            }
-
-            Assert.assertEquals(2, topNArraySize);
-        }
+        assertTopQueriesCount(2, "cpu");
     }
 
     /**
@@ -112,10 +90,10 @@ private String topQueriesByResourceUsagesSettings() {
         return "{\n"
             + "    \"persistent\" : {\n"
             + "        \"search.insights.top_queries.memory.enabled\" : \"true\",\n"
-            + "        \"search.insights.top_queries.memory.window_size\" : \"600s\",\n"
+            + "        \"search.insights.top_queries.memory.window_size\" : \"1m\",\n"
             + "        \"search.insights.top_queries.memory.top_n_size\" : \"5\",\n"
             + "        \"search.insights.top_queries.cpu.enabled\" : \"true\",\n"
-            + "        \"search.insights.top_queries.cpu.window_size\" : \"600s\",\n"
+            + "        \"search.insights.top_queries.cpu.window_size\" : \"1m\",\n"
             + "        \"search.insights.top_queries.cpu.top_n_size\" : 5,\n"
             + "        \"search.insights.top_queries.group_by\" : \"none\"\n"
             + "    }\n"