Fix tests

Signed-off-by: Craig Perkins <[email protected]>
opensearch-project · Jul 27, 2024 · 5be69fc · 5be69fc
1 parent 288d8f2
commit 5be69fc
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 7 deletions.
diff --git a/...r/src/main/java/org/opensearch/action/support/replication/TransportReplicationAction.java b/...r/src/main/java/org/opensearch/action/support/replication/TransportReplicationAction.java
@@ -143,7 +143,7 @@ public abstract class TransportReplicationAction<
     public static final String REPLICA_ACTION_SUFFIX = "[r]";
 
     protected final ThreadPool threadPool;
-    protected final InternalThreadContextWrapper tcWrapper;
+    protected volatile InternalThreadContextWrapper tcWrapper;
     protected final TransportService transportService;
     protected final ClusterService clusterService;
     protected final ShardStateAction shardStateAction;
@@ -243,8 +243,6 @@ protected TransportReplicationAction(
         this.threadPool = threadPool;
         if (threadPool != null) {
             this.tcWrapper = InternalThreadContextWrapper.from(threadPool.getThreadContext());
-        } else {
-            this.tcWrapper = InternalThreadContextWrapper.from(transportService.getThreadPool().getThreadContext());
         }
         this.transportService = transportService;
         this.clusterService = clusterService;

diff --git a/server/src/main/java/org/opensearch/transport/RemoteConnectionStrategy.java b/server/src/main/java/org/opensearch/transport/RemoteConnectionStrategy.java
@@ -172,7 +172,7 @@ public Writeable.Reader<RemoteConnectionInfo.ModeInfo> getReader() {
     ) {
         this.clusterAlias = clusterAlias;
         this.transportService = transportService;
-        this.tcWrapper = InternalThreadContextWrapper.from(transportService.threadPool.getThreadContext());
+        this.tcWrapper = InternalThreadContextWrapper.from(transportService.getThreadPool().getThreadContext());
         this.connectionManager = connectionManager;
         this.maxPendingConnectionListeners = REMOTE_MAX_PENDING_CONNECTION_LISTENERS.get(settings);
         connectionManager.addListener(this);

diff --git a/server/src/test/java/org/opensearch/transport/RemoteConnectionStrategyTests.java b/server/src/test/java/org/opensearch/transport/RemoteConnectionStrategyTests.java
@@ -36,17 +36,21 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.ThreadPool;
 
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class RemoteConnectionStrategyTests extends OpenSearchTestCase {
 
     public void testStrategyChangeMeansThatStrategyMustBeRebuilt() {
         ClusterConnectionManager connectionManager = new ClusterConnectionManager(Settings.EMPTY, mock(Transport.class));
         RemoteConnectionManager remoteConnectionManager = new RemoteConnectionManager("cluster-alias", connectionManager);
+        TransportService mockTransportService = mock(TransportService.class);
+        when(mockTransportService.getThreadPool()).thenReturn(mock(ThreadPool.class));
         FakeConnectionStrategy first = new FakeConnectionStrategy(
             "cluster-alias",
-            mock(TransportService.class),
+            mockTransportService,
             remoteConnectionManager,
             RemoteConnectionStrategy.ConnectionStrategy.PROXY
         );
@@ -60,9 +64,11 @@ public void testStrategyChangeMeansThatStrategyMustBeRebuilt() {
     public void testSameStrategyChangeMeansThatStrategyDoesNotNeedToBeRebuilt() {
         ClusterConnectionManager connectionManager = new ClusterConnectionManager(Settings.EMPTY, mock(Transport.class));
         RemoteConnectionManager remoteConnectionManager = new RemoteConnectionManager("cluster-alias", connectionManager);
+        TransportService mockTransportService = mock(TransportService.class);
+        when(mockTransportService.getThreadPool()).thenReturn(mock(ThreadPool.class));
         FakeConnectionStrategy first = new FakeConnectionStrategy(
             "cluster-alias",
-            mock(TransportService.class),
+            mockTransportService,
             remoteConnectionManager,
             RemoteConnectionStrategy.ConnectionStrategy.PROXY
         );
@@ -78,9 +84,11 @@ public void testChangeInConnectionProfileMeansTheStrategyMustBeRebuilt() {
         assertEquals(TimeValue.MINUS_ONE, connectionManager.getConnectionProfile().getPingInterval());
         assertEquals(false, connectionManager.getConnectionProfile().getCompressionEnabled());
         RemoteConnectionManager remoteConnectionManager = new RemoteConnectionManager("cluster-alias", connectionManager);
+        TransportService mockTransportService = mock(TransportService.class);
+        when(mockTransportService.getThreadPool()).thenReturn(mock(ThreadPool.class));
         FakeConnectionStrategy first = new FakeConnectionStrategy(
             "cluster-alias",
-            mock(TransportService.class),
+            mockTransportService,
             remoteConnectionManager,
             RemoteConnectionStrategy.ConnectionStrategy.PROXY
         );