test changes

Signed-off-by: Gagan Juneja <[email protected]>
Gaganjuneja · Jul 18, 2024 · 161c7f5 · 161c7f5
1 parent ff9d087
commit 161c7f5
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/...rg/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java b/...rg/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java
@@ -59,9 +59,15 @@ public static void setup() {
     @Before
     public void init() {
         initMocks(this);
+        OpenSearchResources.INSTANCE.setMetricsRegistry(metricsRegistry);
         Mockito.when(controller.isPerformanceAnalyzerEnabled()).thenReturn(true);
+        Mockito.when(
+                        metricsRegistry.createHistogram(
+                                Mockito.eq("CPU_Utilization"),
+                                Mockito.anyString(),
+                                Mockito.eq("rate")))
+                .thenReturn(histogram);
         searchListener = new RTFPerformanceAnalyzerSearchListener(controller);
-        OpenSearchResources.INSTANCE.setMetricsRegistry(metricsRegistry);
         assertEquals(
                 RTFPerformanceAnalyzerSearchListener.class.getSimpleName(),
                 searchListener.toString());
@@ -70,16 +76,18 @@ public void init() {
     @Test
     public void tesSearchListener() {
         Mockito.when(controller.getCollectorsSettingValue())
-                .thenReturn(Util.CollectorMode.TELEMETRY.getValue());
+                .thenReturn(Util.CollectorMode.RCA.getValue());
         assertTrue(searchListener.getSearchListener() instanceof NoOpSearchListener);
 
         Mockito.when(controller.getCollectorsSettingValue())
-                .thenReturn(Util.CollectorMode.RCA.getValue());
-        assertTrue(searchListener.getSearchListener() instanceof PerformanceAnalyzerSearchListener);
+                .thenReturn(Util.CollectorMode.TELEMETRY.getValue());
+        assertTrue(
+                searchListener.getSearchListener() instanceof RTFPerformanceAnalyzerSearchListener);
 
         Mockito.when(controller.getCollectorsSettingValue())
                 .thenReturn(Util.CollectorMode.DUAL.getValue());
-        assertTrue(searchListener.getSearchListener() instanceof PerformanceAnalyzerSearchListener);
+        assertTrue(
+                searchListener.getSearchListener() instanceof RTFPerformanceAnalyzerSearchListener);
     }
 
     @Test
@@ -173,6 +181,7 @@ public void testTaskCompletionListener() {
     private void initializeValidSearchContext(boolean isValid) {
         if (isValid) {
             Mockito.when(searchContext.request()).thenReturn(shardSearchRequest);
+            Mockito.when(searchContext.getTask()).thenReturn(task);
             Mockito.when(shardSearchRequest.shardId()).thenReturn(shardId);
         } else {
             Mockito.when(searchContext.request()).thenReturn(null);