Fix tests

Signed-off-by: Craig Perkins <[email protected]>

modules/reindex/src/test/java/org/opensearch/index/reindex/remote/RemoteScrollableHitSourceTests.java

@@ -606,6 +606,8 @@ protected <T> Future<T> doExecute(
                 FutureCallback<T> callback
             ) {
                 try {
+                    // Throw away the current thread context to simulate running async httpclient's thread pool
+                    threadPool.getThreadContext().stashContext();
                     ClassicHttpRequest request = getRequest(requestProducer);
                     URL resource = resources[responseCount];
                     String path = paths[responseCount++];

qa/logging-config/src/test/java/org/opensearch/common/logging/JsonLoggerTests.java

@@ -408,8 +408,11 @@ protected List<String> featureValueOf(JsonLogLine actual) {
 
     private void withThreadContext(CheckedConsumer<ThreadContext, Exception> consumer) throws Exception {
         final ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
-        HeaderWarning.setThreadContext(threadContext);
-        consumer.accept(threadContext);
-        HeaderWarning.removeThreadContext(threadContext);
+        try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
+            HeaderWarning.setThreadContext(threadContext);
+            consumer.accept(threadContext);
+        } finally {
+            HeaderWarning.removeThreadContext(threadContext);
+        }
     }
 }

server/src/main/java/org/opensearch/index/shard/RefreshListeners.java

@@ -99,7 +99,7 @@ public RefreshListeners(
         final IntSupplier getMaxRefreshListeners,
         final Runnable forceRefresh,
         final Logger logger,
-        ThreadContext threadContext,
+        final ThreadContext threadContext,
         final MeanMetric refreshMetric
     ) {
         this.getMaxRefreshListeners = getMaxRefreshListeners;

server/src/main/java/org/opensearch/plugins/TelemetryAwarePlugin.java

@@ -8,7 +8,7 @@
 
 package org.opensearch.plugins;
 
-import org.opensearch.client.node.NodeClient;
+import org.opensearch.client.Client;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.annotation.ExperimentalApi;
@@ -63,7 +63,7 @@ public interface TelemetryAwarePlugin {
      * @param contextSwitcher A client for switching to plugin system context
      */
     default Collection<Object> createComponents(
-        NodeClient client,
+        Client client,
         ClusterService clusterService,
         ThreadPool threadPool,
         ResourceWatcherService resourceWatcherService,

server/src/test/java/org/opensearch/node/NodeTests.java

@@ -34,7 +34,7 @@
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.opensearch.bootstrap.BootstrapCheck;
 import org.opensearch.bootstrap.BootstrapContext;
-import org.opensearch.client.node.NodeClient;
+import org.opensearch.client.Client;
 import org.opensearch.cluster.ClusterName;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.node.DiscoveryNodeRole;
@@ -467,7 +467,7 @@ public MetricsRegistry getMetricsRegistry() {
     public static class MockTelemetryAwarePlugin extends Plugin implements TelemetryAwarePlugin {
         @Override
         public Collection<Object> createComponents(
-            NodeClient client,
+            Client client,
             ClusterService clusterService,
             ThreadPool threadPool,
             ResourceWatcherService resourceWatcherService,

server/src/test/java/org/opensearch/rest/action/admin/indices/RestValidateQueryActionTests.java

@@ -56,6 +56,7 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.usage.UsageService;
 import org.junit.AfterClass;
+import org.junit.Before;
 import org.junit.BeforeClass;
 
 import java.util.Collections;
@@ -112,6 +113,12 @@ protected void doExecute(Task task, ActionRequest request, ActionListener listen
         controller.registerHandler(action);
     }
 
+    @Before
+    public void ensureCleanContext() {
+        // Make sure we have a clean context for each test
+        threadPool.getThreadContext().stashContext();
+    }
+
     @AfterClass
     public static void terminateThreadPool() throws InterruptedException {
         terminate(threadPool);

server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java

@@ -8,8 +8,6 @@
 
 package org.opensearch.telemetry.tracing;
 
-import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
-
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
@@ -21,7 +19,6 @@
 import org.opensearch.telemetry.tracing.noop.NoopTracer;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;
-import org.opensearch.threadpool.ThreadPool;
 import org.junit.After;
 import org.junit.Before;
 
@@ -41,10 +38,8 @@
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.CoreMatchers.nullValue;
 
-@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
 public class ThreadContextBasedTracerContextStorageTests extends OpenSearchTestCase {
     private Tracer tracer;
-    private ThreadPool threadPool;
     private ThreadContext threadContext;
     private TracerContextStorage<String, Span> threadContextStorage;
     private ExecutorService executorService;

test/framework/src/main/java/org/opensearch/test/OpenSearchTestCase.java

@@ -618,7 +618,7 @@ private void assertWarnings(boolean stripXContentPosition, List<String> actualWa
      */
     private void resetDeprecationLogger() {
         // "clear" context by stashing current values and dropping the returned StoredContext
-        // threadContext.stashContext();
+        threadContext.stashContext();
     }
 
     private static final List<StatusData> statusData = new ArrayList<>();