From e6293d32292a8a24400245047626faced2c6eb48 Mon Sep 17 00:00:00 2001
From: Siddhant Deshmukh <deshsid@amazon.com>
Date: Tue, 16 Jul 2024 16:51:09 -0700
Subject: [PATCH] Fix tests

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
---
 .../service/categorizer/SearchQueryCategorizer.java   |  7 +++++++
 .../plugin/insights/QueryInsightsPluginTests.java     | 11 ++++++++++-
 .../core/service/QueryInsightsServiceTests.java       |  2 +-
 .../categorizor/SearchQueryCategorizerTests.java      |  7 +++++++
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java
index 851912a6..b941ed71 100644
--- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java
+++ b/src/main/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
index 0bb43da0..443ca8b8 100644
--- 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
index 60ed6529..06702e81 100644
--- 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/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java
index e511ef3b..643b5141 100644
--- a/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java
+++ b/src/test/java/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(