From 161c7f5f4db01cbbfa5e867c6cd3b78823019e0a Mon Sep 17 00:00:00 2001
From: Gagan Juneja <gjjuneja@amazon.com>
Date: Thu, 18 Jul 2024 06:26:38 +0530
Subject: [PATCH] test changes

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
---
 ...erformanceAnalyzerSearchListenerTests.java | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/test/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java b/src/test/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java
index ae7b496b..87b312a6 100644
--- a/src/test/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java
+++ b/src/test/java/org/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);