Fix test case

Signed-off-by: Gagan Juneja <[email protected]>
Gaganjuneja · Sep 8, 2023 · eb99f2e · eb99f2e
1 parent e7cf6ba
commit eb99f2e
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/server/src/main/java/org/opensearch/transport/TransportService.java b/server/src/main/java/org/opensearch/transport/TransportService.java
@@ -1043,6 +1043,7 @@ public void onRejection(Exception e) {
                     @Override
                     public void onFailure(Exception e) {
                         logger.error("Failing from onFailure ", e);
+                        contextToNotify.handler().handleRejection(e);
                         logger.warn(
                             () -> new ParameterizedMessage(
                                 "failed to notify response handler on exception, action: {}",

diff --git a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockTracingTelemetry.java b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockTracingTelemetry.java
@@ -13,12 +13,15 @@
 import org.opensearch.telemetry.tracing.TracingTelemetry;
 import org.opensearch.telemetry.tracing.attributes.Attributes;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 /**
  * Mock {@link TracingTelemetry} implementation for testing.
  */
 public class MockTracingTelemetry implements TracingTelemetry {
 
     private final SpanProcessor spanProcessor = new StrictCheckSpanProcessor();
+    private final AtomicBoolean shutdown = new AtomicBoolean();
 
     /**
      * Base constructor.
@@ -28,7 +31,9 @@ public MockTracingTelemetry() {}
     @Override
     public Span createSpan(String spanName, Span parentSpan, Attributes attributes) {
         Span span = new MockSpan(spanName, parentSpan, spanProcessor, attributes);
-        spanProcessor.onStart(span);
+        if (shutdown.get() == false) {
+            spanProcessor.onStart(span);
+        }
         return span;
     }
 
@@ -39,7 +44,8 @@ public TracingContextPropagator getContextPropagator() {
 
     @Override
     public void close() {
-        StrictCheckSpanProcessor.shutdown();
+        // StrictCheckSpanProcessor.shutdown();
+        shutdown.set(true);
     }
 
 }
diff --git a/...lemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java b/...lemetry/src/main/java/org/opensearch/test/telemetry/tracing/StrictCheckSpanProcessor.java
@@ -31,7 +31,7 @@ public StrictCheckSpanProcessor() {
     }
 
     private void initProcessor() {
-        restartTheProcessor();
+        // restartTheProcessor();
     }
 
     private static Map<String, MockSpanData> spanMap = new ConcurrentHashMap<>();
@@ -81,7 +81,7 @@ private MockSpanData toMockSpanData(Span span) {
      */
     public static synchronized void shutdown() {
         if (shutdown.get() == false) {
-            shutdown.set(true);
+            // shutdown.set(true);
         }
     }
 
@@ -109,7 +109,7 @@ public static void validateTracingStateOnShutdown() {
                 spanMap.clear();
                 throw e;
             } finally {
-                restartTheProcessor();
+                // restartTheProcessor();
             }
         }