From 518b39f167c9e1ad3405d0d5d70b750227b50dab Mon Sep 17 00:00:00 2001
From: Gagan Juneja <gjjuneja@amazon.com>
Date: Fri, 8 Sep 2023 16:43:11 +0530
Subject: [PATCH] Fix test case

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
---
 .../org/opensearch/search/aggregations/bucket/IpRangeIT.java | 5 +++++
 .../java/org/opensearch/test/OpenSearchIntegTestCase.java    | 1 +
 .../test/telemetry/tracing/StrictCheckSpanProcessor.java     | 4 +---
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/IpRangeIT.java b/server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/IpRangeIT.java
index f8f666aaa3c1b..eada2bafb6258 100644
--- a/server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/IpRangeIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/IpRangeIT.java
@@ -56,6 +56,11 @@
 @OpenSearchIntegTestCase.SuiteScopeTestCase
 public class IpRangeIT extends OpenSearchIntegTestCase {
 
+    @Override
+    protected boolean addMockTelemetryPlugin() {
+        return true;
+    }
+
     public static class DummyScriptPlugin extends MockScriptPlugin {
         @Override
         public Map<String, Function<Map<String, Object>, Object>> pluginScripts() {
diff --git a/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java b/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java
index da89d3424aef8..abc0e90bde3ee 100644
--- a/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java
+++ b/test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java
@@ -2328,6 +2328,7 @@ private static void initializeSuiteScope() throws Exception {
                 INSTANCE.beforeInternal();
                 INSTANCE.setupSuiteScopeCluster();
                 success = true;
+                StrictCheckSpanProcessor.restartTheProcessor();
             } finally {
                 if (!success) {
                     afterClass();
diff --git a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java
index 58ec6897cce9a..5d74d6b97ee0b 100644
--- a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java
+++ b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java
@@ -97,9 +97,7 @@ public static synchronized void shutdown() {
      * restartTheProcessor
      */
     public static synchronized void restartTheProcessor() {
-        if (shutdown.get() == true) {
-            shutdown.set(false);
-        }
+        spanMap.clear();
     }
 
     /**