Fix tests

Signed-off-by: Siddhant Deshmukh <[email protected]>
opensearch-project · Jul 16, 2024 · e6293d3 · e6293d3
1 parent 6b925be
commit e6293d3
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 2 deletions.
diff --git a/.../java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java b/.../java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java
@@ -129,4 +129,11 @@ private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
     public SearchQueryCounters getSearchQueryCounters() {
         return this.searchQueryCounters;
     }
+
+    /**
+     * Reset the search query categorizer and its counters
+     */
+    public void reset() {
+        instance = null;
+    }
 }
diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java
@@ -23,6 +23,8 @@
 import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.rest.RestHandler;
+import org.opensearch.telemetry.metrics.Counter;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
 import org.opensearch.test.ClusterServiceUtils;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ExecutorBuilder;
@@ -32,7 +34,9 @@
 import java.util.Arrays;
 import java.util.List;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class QueryInsightsPluginTests extends OpenSearchTestCase {
 
@@ -41,6 +45,7 @@ public class QueryInsightsPluginTests extends OpenSearchTestCase {
     private final Client client = mock(Client.class);
     private ClusterService clusterService;
     private final ThreadPool threadPool = mock(ThreadPool.class);
+    private MetricsRegistry metricsRegistry = mock(MetricsRegistry.class);
 
     @Before
     public void setup() {
@@ -50,6 +55,10 @@ public void setup() {
         ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
         QueryInsightsTestUtils.registerAllQueryInsightsSettings(clusterSettings);
         clusterService = ClusterServiceUtils.createClusterService(settings, clusterSettings, threadPool);
+
+        when(metricsRegistry.createCounter(any(String.class), any(String.class), any(String.class))).thenAnswer(
+            invocation -> mock(Counter.class)
+        );
     }
 
     public void testGetSettings() {
@@ -87,7 +96,7 @@ public void testCreateComponent() {
             null,
             null,
             null,
-            null
+            metricsRegistry
         );
         assertEquals(2, components.size());
         assertTrue(components.get(0) instanceof QueryInsightsService);

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
@@ -67,7 +67,7 @@ public void testClose() {
     public void testSearchQueryMetricsEnabled() {
         // Initially, searchQueryMetricsEnabled should be false and searchQueryCategorizer should be null
         assertFalse(queryInsightsService.isSearchQueryMetricsEnabled());
-        assertNull(queryInsightsService.getSearchQueryCategorizer());
+        assertNotNull(queryInsightsService.getSearchQueryCategorizer());
 
         // Enable search query metrics
         queryInsightsService.setSearchQueryMetricsEnabled(true);

diff --git a/.../org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java b/.../org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java
@@ -8,6 +8,8 @@
 
 package org.opensearch.plugin.insights.core.service.categorizor;
 
+import org.junit.After;
+import org.mockito.Mockito;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.BoostingQueryBuilder;
 import org.opensearch.index.query.MatchNoneQueryBuilder;
@@ -61,6 +63,11 @@ public void setup() {
         searchQueryCategorizer = SearchQueryCategorizer.getInstance(metricsRegistry);
     }
 
+    @After
+    public void cleanup() {
+        searchQueryCategorizer.reset();
+    }
+
     public void testAggregationsQuery() {
         SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
         sourceBuilder.aggregation(