Avoid getters that are only used in tests

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
opensearch-project · Jul 18, 2024 · f724ef5 · f724ef5
1 parent 2486d24
commit f724ef5
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 31 deletions.
diff --git a/...ain/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java b/...ain/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java
@@ -171,28 +171,4 @@ private Counter createQueryCounter(String counterName) {
         );
         return counter;
     }
-
-    /**
-     * Get Query type latency histogram
-     * @return histogram
-     */
-    public Histogram getQueryTypeLatencyHistogram() {
-        return queryTypeLatencyHistogram;
-    }
-
-    /**
-     * Get Query type cpu histogram
-     * @return histogram
-     */
-    public Histogram getQueryTypeCpuHistogram() {
-        return queryTypeCpuHistogram;
-    }
-
-    /**
-     * Get Query type memory histogram
-     * @return histogram
-     */
-    public Histogram getQueryTypeMemoryHistogram() {
-        return queryTypeMemoryHistogram;
-    }
 }
diff --git a/.../org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java b/.../org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java
@@ -9,7 +9,8 @@
 package org.opensearch.plugin.insights.core.service.categorizor;
 
 import org.junit.After;
-import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.BoostingQueryBuilder;
 import org.opensearch.index.query.MatchNoneQueryBuilder;
@@ -39,6 +40,7 @@
 import org.junit.Before;
 
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Map;
 
 import org.mockito.ArgumentCaptor;
@@ -59,16 +61,29 @@ public final class SearchQueryCategorizerTests extends OpenSearchTestCase {
 
     private SearchQueryCategorizer searchQueryCategorizer;
 
+    private Map<String, Histogram> histogramMap = new HashMap<>();
+
     @Before
     public void setup() {
         SearchQueryCategorizer.getInstance(mock(MetricsRegistry.class)).reset();
         metricsRegistry = mock(MetricsRegistry.class);
         when(metricsRegistry.createCounter(any(String.class), any(String.class), any(String.class))).thenAnswer(
             invocation -> mock(Counter.class)
         );
-        when(metricsRegistry.createHistogram(any(String.class), any(String.class), any(String.class))).thenAnswer(
-            invocation -> mock(Histogram.class)
-        );
+
+        when(metricsRegistry.createHistogram(any(String.class), any(String.class), any(String.class)))
+            .thenAnswer(new Answer<Histogram>() {
+                @Override
+                public Histogram answer(InvocationOnMock invocation) throws Throwable {
+                    // Extract arguments to identify which histogram is being created
+                    String name = invocation.getArgument(0);
+                    // Create a mock histogram
+                    Histogram histogram = mock(Histogram.class);
+                    // Store histogram in map for lookup
+                    histogramMap.put(name, histogram);
+                    return histogram;
+                }
+            });
         searchQueryCategorizer = SearchQueryCategorizer.getInstance(metricsRegistry);
     }
 
@@ -308,8 +323,13 @@ private void verifyMeasurementHistogramsIncremented(SearchQueryRecord record, in
         Double expectedCpu = record.getMeasurement(MetricType.CPU).doubleValue();
         Double expectedMemory = record.getMeasurement(MetricType.MEMORY).doubleValue();
 
-        verify(searchQueryCategorizer.getSearchQueryCounters().getQueryTypeLatencyHistogram(), times(times)).record(eq(expectedLatency), any(Tags.class));
-        verify(searchQueryCategorizer.getSearchQueryCounters().getQueryTypeCpuHistogram(), times(times)).record(eq(expectedCpu), any(Tags.class));
-        verify(searchQueryCategorizer.getSearchQueryCounters().getQueryTypeMemoryHistogram(), times(times)).record(eq(expectedMemory), any(Tags.class));
+        Histogram queryTypeLatencyHistogram = histogramMap.get("search.query.type.latency.histogram");
+        Histogram queryTypeCpuHistogram = histogramMap.get("search.query.type.cpu.histogram");
+        Histogram queryTypeMemoryHistogram = histogramMap.get("search.query.type.memory.histogram");
+
+
+        verify(queryTypeLatencyHistogram, times(times)).record(eq(expectedLatency), any(Tags.class));
+        verify(queryTypeCpuHistogram, times(times)).record(eq(expectedCpu), any(Tags.class));
+        verify(queryTypeMemoryHistogram, times(times)).record(eq(expectedMemory), any(Tags.class));
     }
 }