Separate mock task manager from mock transport service (#138609)

* Separate mock task manager from mock transport service

* Remove MockTransportService dependency from MockNode

* Revert "Remove MockTransportService dependency from MockNode"

This reverts commit 11ddaaf.

* Update test/framework/src/main/java/org/elasticsearch/node/MockNode.java

Co-authored-by: Yang Wang <yang.wang@elastic.co>

---------

Co-authored-by: Yang Wang <yang.wang@elastic.co>

Author: jxie-1

server/src/main/java/org/elasticsearch/node/NodeConstruction.java

@@ -711,7 +711,8 @@ private void construct(
         modules.bindToInstance(RootObjectMapperNamespaceValidator.class, namespaceValidator);
 
         assert nodeEnvironment.nodeId() != null : "node ID must be set before constructing the Node";
-        TaskManager taskManager = new TaskManager(
+        TaskManager taskManager = serviceProvider.newTaskManager(
+            pluginsService,
             settings,
             threadPool,
             Stream.concat(

server/src/main/java/org/elasticsearch/node/NodeServiceProvider.java

@@ -50,6 +50,7 @@
 import org.elasticsearch.transport.TransportService;
 
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.function.LongSupplier;
 
@@ -63,6 +64,17 @@ PluginsService newPluginService(Environment initialEnvironment, PluginsLoader pl
         return new PluginsService(initialEnvironment.settings(), initialEnvironment.configDir(), pluginsLoader);
     }
 
+    TaskManager newTaskManager(
+        PluginsService pluginsService,
+        Settings settings,
+        ThreadPool threadPool,
+        Set<String> taskHeaders,
+        Tracer tracer,
+        String nodeId
+    ) {
+        return new TaskManager(settings, threadPool, taskHeaders, tracer, nodeId);
+    }
+
     ScriptService newScriptService(
         PluginsService pluginsService,
         Settings settings,

server/src/test/java/org/elasticsearch/action/admin/cluster/node/tasks/TaskManagerTestCase.java

@@ -62,7 +62,6 @@
 import static java.util.Collections.emptySet;
 import static org.elasticsearch.test.ClusterServiceUtils.createClusterService;
 import static org.elasticsearch.test.ClusterServiceUtils.setState;
-import static org.elasticsearch.test.transport.MockTransportService.createTaskManager;
 
 /**
  * The test case for unit testing task manager and related transport actions
@@ -177,7 +176,7 @@ public TestNode(String name, ThreadPool threadPool, Settings settings) {
                 discoveryNode.set(DiscoveryNodeUtils.create(name, address.publishAddress(), emptyMap(), emptySet()));
                 return discoveryNode.get();
             };
-            TaskManager taskManager = createTaskManager(settings, threadPool, emptySet(), Tracer.NOOP, name);
+            TaskManager taskManager = MockTaskManager.create(settings, threadPool, emptySet(), Tracer.NOOP, name);
             transportService = new TransportService(
                 settings,
                 new Netty4Transport(

test/framework/src/main/java/org/elasticsearch/node/MockNode.java

@@ -49,6 +49,7 @@
 import org.elasticsearch.telemetry.tracing.Tracer;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.MockHttpTransport;
+import org.elasticsearch.test.tasks.MockTaskManager;
 import org.elasticsearch.test.transport.MockTransportService;
 import org.elasticsearch.test.transport.StubbableTransport;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -63,6 +64,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.function.LongSupplier;
@@ -78,6 +80,21 @@
 public class MockNode extends Node {
 
     private static class MockServiceProvider extends NodeServiceProvider {
+        @Override
+        TaskManager newTaskManager(
+            PluginsService pluginsService,
+            Settings settings,
+            ThreadPool threadPool,
+            Set<String> taskHeaders,
+            Tracer tracer,
+            String nodeId
+        ) {
+            if (pluginsService.filterPlugins(MockTransportService.TestPlugin.class).findAny().isEmpty()) {
+                return super.newTaskManager(pluginsService, settings, threadPool, taskHeaders, tracer, nodeId);
+            }
+            return MockTaskManager.create(settings, threadPool, taskHeaders, tracer, nodeId);
+        }
+
         @Override
         BigArrays newBigArrays(
             PluginsService pluginsService,
@@ -208,7 +225,7 @@ protected TransportService newTransportService(
                     interceptor,
                     localNodeFactory,
                     clusterSettings,
-                    MockTransportService.createTaskManager(settings, threadPool, taskManager.getTaskHeaders(), Tracer.NOOP, nodeId),
+                    taskManager,
                     linkedProjectConfigService,
                     telemetryProvider,
                     projectResolver

test/framework/src/main/java/org/elasticsearch/test/tasks/MockTaskManager.java

@@ -26,6 +26,7 @@
 import java.util.concurrent.CopyOnWriteArrayList;
 
 import static org.elasticsearch.core.Strings.format;
+import static org.mockito.Mockito.spy;
 
 /**
  * A mock task manager that allows adding listeners for events
@@ -52,6 +53,13 @@ public MockTaskManager(Settings settings, ThreadPool threadPool, Set<String> tas
         super(settings, threadPool, taskHeaders, Tracer.NOOP);
     }
 
+    public static TaskManager create(Settings settings, ThreadPool threadPool, Set<String> taskHeaders, Tracer tracer, String nodeId) {
+        var taskManager = MockTaskManager.USE_MOCK_TASK_MANAGER_SETTING.get(settings)
+            ? new MockTaskManager(settings, threadPool, taskHeaders)
+            : new TaskManager(settings, threadPool, taskHeaders, tracer, nodeId);
+        return MockTaskManager.SPY_TASK_MANAGER_SETTING.get(settings) ? spy(taskManager) : taskManager;
+    }
+
     @Override
     public Task register(String type, String action, TaskAwareRequest request) {
         Task task = super.register(type, action, request);

test/framework/src/main/java/org/elasticsearch/test/transport/MockTransportService.java

@@ -94,7 +94,6 @@
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.spy;
 
 /**
  * A mock delegate service that allows to simulate different network topology failures.
@@ -265,7 +264,7 @@ public MockTransportService(
             interceptor,
             localNodeFactory,
             clusterSettings,
-            createTaskManager(settings, threadPool, taskHeaders, Tracer.NOOP, nodeId),
+            MockTaskManager.create(settings, threadPool, taskHeaders, Tracer.NOOP, nodeId),
             new ClusterSettingsLinkedProjectConfigService(settings, clusterSettings, DefaultProjectResolver.INSTANCE),
             new TelemetryProvider() {
                 final MeterRegistry meterRegistry = new RecordingMeterRegistry();
@@ -330,31 +329,6 @@ private static TransportAddress[] extractTransportAddresses(TransportService tra
         return transportAddresses.toArray(new TransportAddress[transportAddresses.size()]);
     }
 
-    public static TaskManager createTaskManager(
-        Settings settings,
-        ThreadPool threadPool,
-        Set<String> taskHeaders,
-        Tracer tracer,
-        String nodeId
-    ) {
-        TaskManager mockTaskManager = createMockTaskManager(settings, threadPool, taskHeaders, tracer, nodeId);
-        return MockTaskManager.SPY_TASK_MANAGER_SETTING.get(settings) ? spy(mockTaskManager) : mockTaskManager;
-    }
-
-    private static TaskManager createMockTaskManager(
-        Settings settings,
-        ThreadPool threadPool,
-        Set<String> taskHeaders,
-        Tracer tracer,
-        String nodeId
-    ) {
-        if (MockTaskManager.USE_MOCK_TASK_MANAGER_SETTING.get(settings)) {
-            return new MockTaskManager(settings, threadPool, taskHeaders);
-        } else {
-            return new TaskManager(settings, threadPool, taskHeaders, tracer, nodeId);
-        }
-    }
-
     /**
      * Clears all the registered rules.
      */