Add unit test and change serialization

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
deshsidd · Jul 2, 2024 · d419ca5 · d419ca5
1 parent eb8ebd8
commit d419ca5
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 13 deletions.
diff --git a/...ights/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java b/...ights/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java
@@ -122,13 +122,6 @@ public QueryInsightsService(
         this.metricsRegistry = metricsRegistry;
     }
 
-    private void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) {
-        this.searchQueryMetricsEnabled = searchQueryMetricsEnabled;
-        if ((this.searchQueryMetricsEnabled == true) && this.searchQueryCategorizer == null) {
-            this.searchQueryCategorizer = new SearchQueryCategorizer(metricsRegistry);
-        }
-    }
-
     /**
      * Ingest the query data into in-memory stores
      *
@@ -175,7 +168,6 @@ public void drainRecords() {
                 logger.error("Error while trying to categorize the queries.", e);
             }
         }
-
     }
 
     /**
@@ -282,6 +274,25 @@ public void setExporter(final MetricType type, final Settings settings) {
         }
     }
 
+    /**
+     * Set search query metrics enabled to enable collection of search query categorization metrics
+     * @param searchQueryMetricsEnabled boolean flag
+     */
+    public void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) {
+        this.searchQueryMetricsEnabled = searchQueryMetricsEnabled;
+        if ((this.searchQueryMetricsEnabled == true) && this.searchQueryCategorizer == null) {
+            this.searchQueryCategorizer = new SearchQueryCategorizer(metricsRegistry);
+        }
+    }
+
+    public boolean isSearchQueryMetricsEnabled() {
+        return this.searchQueryMetricsEnabled;
+    }
+
+    public SearchQueryCategorizer getSearchQueryCategorizer() {
+        return this.searchQueryCategorizer;
+    }
+
     /**
      * Validate the exporter config for a metricType
      *

diff --git a/...-insights/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java b/...-insights/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java
@@ -134,11 +134,7 @@ public XContentBuilder toXContent(final XContentBuilder builder, final ToXConten
     public void writeTo(final StreamOutput out) throws IOException {
         out.writeLong(timestamp);
         out.writeMap(measurements, (stream, metricType) -> MetricType.writeTo(out, metricType), StreamOutput::writeGenericValue);
-        out.writeMap(
-            attributes,
-            (stream, attribute) -> Attribute.writeTo(out, attribute),
-            (stream, attributeValue) -> Attribute.writeValueTo(out, attributeValue)
-        );
+        out.writeMap(attributes, (stream, attribute) -> Attribute.writeTo(out, attribute), StreamOutput::writeGenericValue);
     }
 
     /**

diff --git a/.../src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java b/.../src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java
@@ -63,4 +63,25 @@ public void testClose() {
             fail("No exception expected when closing query insights service");
         }
     }
+
+    public void testSearchQueryMetricsEnabled() {
+        // Initially, searchQueryMetricsEnabled should be false and searchQueryCategorizer should be null
+        assertFalse(queryInsightsService.isSearchQueryMetricsEnabled());
+        assertNull(queryInsightsService.getSearchQueryCategorizer());
+
+        // Enable search query metrics
+        queryInsightsService.setSearchQueryMetricsEnabled(true);
+
+        // Assert that searchQueryMetricsEnabled is true and searchQueryCategorizer is initialized
+        assertTrue(queryInsightsService.isSearchQueryMetricsEnabled());
+        assertNotNull(queryInsightsService.getSearchQueryCategorizer());
+
+        // Disable search query metrics
+        queryInsightsService.setSearchQueryMetricsEnabled(false);
+
+        // Assert that searchQueryMetricsEnabled is false and searchQueryCategorizer is not null
+        assertFalse(queryInsightsService.isSearchQueryMetricsEnabled());
+        assertNotNull(queryInsightsService.getSearchQueryCategorizer());
+
+    }
 }