Refactor test code to reduce boilerplate

Author: surya-estc

server/src/test/java/org/elasticsearch/cluster/InternalClusterInfoServiceRefreshTests.java

@@ -46,19 +46,12 @@
 public class InternalClusterInfoServiceRefreshTests extends ESTestCase {
 
     public void testEstimatedHeapUsageCollectorSuccessAndFailure() {
-        final Settings settings = Settings.builder()
-            .put(DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.getKey(), false)
-            .put(
-                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
-                WriteLoadConstraintSettings.WriteLoadDeciderStatus.DISABLED
-            )
-            .put(InternalClusterInfoService.CLUSTER_ROUTING_ALLOCATION_ESTIMATED_HEAP_THRESHOLD_DECIDER_ENABLED.getKey(), true)
-            .build();
-        final ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
-        final DeterministicTaskQueue deterministicTaskQueue = new DeterministicTaskQueue();
-        final ThreadPool threadPool = deterministicTaskQueue.getThreadPool();
+        final Settings settings = baseSettingsBuilder().put(
+            InternalClusterInfoService.CLUSTER_ROUTING_ALLOCATION_ESTIMATED_HEAP_THRESHOLD_DECIDER_ENABLED.getKey(),
+            true
+        ).build();
 
-        try (ClusterService clusterService = ClusterServiceUtils.createClusterService(threadPool, clusterSettings)) {
+        try (RefreshTestContext context = RefreshTestContext.create(settings)) {
             final Map<String, NodeHeapEstimates> nodeHeapEstimates = Map.of("node-id", new NodeHeapEstimates(100L, 20L));
             final ShardId shardId = new ShardId("index", "uuid", 0);
             final ShardHeapUsageEstimates shardHeapUsageEstimates = new ShardHeapUsageEstimates(
@@ -78,43 +71,13 @@ public void testEstimatedHeapUsageCollectorSuccessAndFailure() {
                 return null;
             }).when(estimatedHeapUsageCollector).collectEstimatedHeapUsage(any());
 
-            final InternalClusterInfoService clusterInfoService = new InternalClusterInfoService(
-                settings,
-                new WriteLoadConstraintSettings(clusterService.getClusterSettings()),
-                clusterService,
-                threadPool,
-                new NoOpClient(threadPool) {
-                    @Override
-                    @SuppressWarnings("unchecked")
-                    protected <Request extends ActionRequest, Response extends ActionResponse> void doExecute(
-                        ActionType<Response> action,
-                        Request request,
-                        ActionListener<Response> listener
-                    ) {
-                        if (request instanceof NodesStatsRequest) {
-                            NodeStats nodeStats = mock(NodeStats.class);
-                            JvmStats jvmStats = mock(JvmStats.class);
-                            JvmStats.Mem mem = mock(JvmStats.Mem.class);
-                            Mockito.when(nodeStats.getNode()).thenReturn(DiscoveryNodeUtils.create("node-id"));
-                            Mockito.when(nodeStats.getJvm()).thenReturn(jvmStats);
-                            Mockito.when(jvmStats.getMem()).thenReturn(mem);
-                            Mockito.when(mem.getHeapMax()).thenReturn(ByteSizeValue.ofBytes(1_000L));
-                            listener.onResponse(
-                                (Response) new NodesStatsResponse(new ClusterName("cluster"), List.of(nodeStats), List.of())
-                            );
-                        } else {
-                            fail("unexpected action: " + action.name());
-                        }
-                    }
-                },
+            final InternalClusterInfoService clusterInfoService = context.createClusterInfoService(
+                nodeStatsClient(context.threadPool()),
                 estimatedHeapUsageCollector,
                 CacheSizesAndCommitmentCollector.EMPTY,
                 PartitionSizeCollector.EMPTY,
-                NodeUsageStatsForThreadPoolsCollector.EMPTY,
                 SearchLaneRequirementsCollector.EMPTY
             );
-            // AsyncRefresh asserts that each refresh notifies at least one registered cluster info listener.
-            clusterInfoService.addListener(ignored -> {});
 
             ClusterInfo clusterInfo = refresh(clusterInfoService);
             verify(estimatedHeapUsageCollector).collectEstimatedHeapUsage(any());
@@ -136,18 +99,7 @@ protected <Request extends ActionRequest, Response extends ActionResponse> void
     }
 
     public void testPartitionSizeCollectorSuccessAndFailure() {
-        final Settings settings = Settings.builder()
-            .put(DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.getKey(), false)
-            .put(
-                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
-                WriteLoadConstraintSettings.WriteLoadDeciderStatus.DISABLED
-            )
-            .build();
-        final ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
-        final DeterministicTaskQueue deterministicTaskQueue = new DeterministicTaskQueue();
-        final ThreadPool threadPool = deterministicTaskQueue.getThreadPool();
-
-        try (ClusterService clusterService = ClusterServiceUtils.createClusterService(threadPool, clusterSettings)) {
+        try (RefreshTestContext context = RefreshTestContext.create(baseSettings())) {
             final Map<String, Long> partitionSizes = Map.of("node-id", 1234L);
             final AtomicBoolean failPartitionSizes = new AtomicBoolean();
             final PartitionSizeCollector partitionSizeCollector = mock(PartitionSizeCollector.class);
@@ -161,20 +113,12 @@ public void testPartitionSizeCollectorSuccessAndFailure() {
                 return null;
             }).when(partitionSizeCollector).collectHostedShardsPartitionSizes(any(), any());
 
-            final InternalClusterInfoService clusterInfoService = new InternalClusterInfoService(
-                settings,
-                new WriteLoadConstraintSettings(clusterService.getClusterSettings()),
-                clusterService,
-                threadPool,
-                new NoOpClient(threadPool),
+            final InternalClusterInfoService clusterInfoService = context.createClusterInfoService(
                 EstimatedHeapUsageCollector.EMPTY,
                 CacheSizesAndCommitmentCollector.EMPTY,
                 partitionSizeCollector,
-                NodeUsageStatsForThreadPoolsCollector.EMPTY,
                 SearchLaneRequirementsCollector.EMPTY
             );
-            // Refresh is a no-op if there are no listeners
-            clusterInfoService.addListener(ignored -> {});
 
             // Success populates the ClusterInfo
             ClusterInfo clusterInfo = refresh(clusterInfoService);
@@ -191,18 +135,7 @@ public void testPartitionSizeCollectorSuccessAndFailure() {
     }
 
     public void testCacheSizesAndCommitmentCollectorSuccessAndFailure() {
-        final Settings settings = Settings.builder()
-            .put(DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.getKey(), false)
-            .put(
-                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
-                WriteLoadConstraintSettings.WriteLoadDeciderStatus.DISABLED
-            )
-            .build();
-        final ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
-        final DeterministicTaskQueue deterministicTaskQueue = new DeterministicTaskQueue();
-        final ThreadPool threadPool = deterministicTaskQueue.getThreadPool();
-
-        try (ClusterService clusterService = ClusterServiceUtils.createClusterService(threadPool, clusterSettings)) {
+        try (RefreshTestContext context = RefreshTestContext.create(baseSettings())) {
             final Map<ShardId, BoostedAndUnboostedCacheRequirements> shardCacheRequirements = Map.of(
                 new ShardId("index", "uuid", 0),
                 new BoostedAndUnboostedCacheRequirements(10L, 20L)
@@ -227,20 +160,12 @@ public void testCacheSizesAndCommitmentCollectorSuccessAndFailure() {
                 return null;
             }).when(cacheSizesAndCommitmentCollector).collectCacheSizesAndCommitmentStats(any(), any());
 
-            final InternalClusterInfoService clusterInfoService = new InternalClusterInfoService(
-                settings,
-                new WriteLoadConstraintSettings(clusterService.getClusterSettings()),
-                clusterService,
-                threadPool,
-                new NoOpClient(threadPool),
+            final InternalClusterInfoService clusterInfoService = context.createClusterInfoService(
                 EstimatedHeapUsageCollector.EMPTY,
                 cacheSizesAndCommitmentCollector,
                 PartitionSizeCollector.EMPTY,
-                NodeUsageStatsForThreadPoolsCollector.EMPTY,
                 SearchLaneRequirementsCollector.EMPTY
             );
-            // AsyncRefresh asserts that each refresh notifies at least one registered cluster info listener.
-            clusterInfoService.addListener(ignored -> {});
 
             ClusterInfo clusterInfo = refresh(clusterInfoService);
             verify(cacheSizesAndCommitmentCollector).collectCacheSizesAndCommitmentStats(any(), any());
@@ -257,18 +182,7 @@ public void testCacheSizesAndCommitmentCollectorSuccessAndFailure() {
     }
 
     public void testSearchLaneRequirementsCollectorSuccessAndFailure() {
-        final Settings settings = Settings.builder()
-            .put(DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.getKey(), false)
-            .put(
-                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
-                WriteLoadConstraintSettings.WriteLoadDeciderStatus.DISABLED
-            )
-            .build();
-        final ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
-        final DeterministicTaskQueue deterministicTaskQueue = new DeterministicTaskQueue();
-        final ThreadPool threadPool = deterministicTaskQueue.getThreadPool();
-
-        try (ClusterService clusterService = ClusterServiceUtils.createClusterService(threadPool, clusterSettings)) {
+        try (RefreshTestContext context = RefreshTestContext.create(baseSettings())) {
             final Map<ShardId, Double> laneRequirements = Map.of(new ShardId("index", "uuid", 0), 3.2);
             final AtomicBoolean failLaneRequirements = new AtomicBoolean();
             final SearchLaneRequirementsCollector searchLaneRequirementsCollector = mock(SearchLaneRequirementsCollector.class);
@@ -282,20 +196,12 @@ public void testSearchLaneRequirementsCollectorSuccessAndFailure() {
                 return null;
             }).when(searchLaneRequirementsCollector).collectSearchLaneRequirements(any(), any());
 
-            final InternalClusterInfoService clusterInfoService = new InternalClusterInfoService(
-                settings,
-                new WriteLoadConstraintSettings(clusterService.getClusterSettings()),
-                clusterService,
-                threadPool,
-                new NoOpClient(threadPool),
+            final InternalClusterInfoService clusterInfoService = context.createClusterInfoService(
                 EstimatedHeapUsageCollector.EMPTY,
                 CacheSizesAndCommitmentCollector.EMPTY,
                 PartitionSizeCollector.EMPTY,
-                NodeUsageStatsForThreadPoolsCollector.EMPTY,
                 searchLaneRequirementsCollector
             );
-            // Refresh is a no-op if there are no listeners
-            clusterInfoService.addListener(ignored -> {});
 
             // Success populates the ClusterInfo
             ClusterInfo clusterInfo = refresh(clusterInfoService);
@@ -310,4 +216,96 @@ public void testSearchLaneRequirementsCollectorSuccessAndFailure() {
             assertThat(clusterInfo.getShardSearchLaneRequirements(), equalTo(Map.of()));
         }
     }
+
+    private static Settings baseSettings() {
+        return baseSettingsBuilder().build();
+    }
+
+    private static Settings.Builder baseSettingsBuilder() {
+        return Settings.builder()
+            .put(DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.getKey(), false)
+            .put(
+                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
+                WriteLoadConstraintSettings.WriteLoadDeciderStatus.DISABLED
+            );
+    }
+
+    private static NoOpClient nodeStatsClient(ThreadPool threadPool) {
+        return new NoOpClient(threadPool) {
+            @Override
+            @SuppressWarnings("unchecked")
+            protected <Request extends ActionRequest, Response extends ActionResponse> void doExecute(
+                ActionType<Response> action,
+                Request request,
+                ActionListener<Response> listener
+            ) {
+                if (request instanceof NodesStatsRequest) {
+                    NodeStats nodeStats = mock(NodeStats.class);
+                    JvmStats jvmStats = mock(JvmStats.class);
+                    JvmStats.Mem mem = mock(JvmStats.Mem.class);
+                    Mockito.when(nodeStats.getNode()).thenReturn(DiscoveryNodeUtils.create("node-id"));
+                    Mockito.when(nodeStats.getJvm()).thenReturn(jvmStats);
+                    Mockito.when(jvmStats.getMem()).thenReturn(mem);
+                    Mockito.when(mem.getHeapMax()).thenReturn(ByteSizeValue.ofBytes(1_000L));
+                    listener.onResponse((Response) new NodesStatsResponse(new ClusterName("cluster"), List.of(nodeStats), List.of()));
+                } else {
+                    fail("unexpected action: " + action.name());
+                }
+            }
+        };
+    }
+
+    private record RefreshTestContext(Settings settings, ThreadPool threadPool, ClusterService clusterService) implements AutoCloseable {
+
+        static RefreshTestContext create(Settings settings) {
+            final ClusterSettings clusterSettings = new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
+            final DeterministicTaskQueue deterministicTaskQueue = new DeterministicTaskQueue();
+            final ThreadPool threadPool = deterministicTaskQueue.getThreadPool();
+            return new RefreshTestContext(settings, threadPool, ClusterServiceUtils.createClusterService(threadPool, clusterSettings));
+        }
+
+        InternalClusterInfoService createClusterInfoService(
+            EstimatedHeapUsageCollector estimatedHeapUsageCollector,
+            CacheSizesAndCommitmentCollector cacheSizesAndCommitmentCollector,
+            PartitionSizeCollector partitionSizeCollector,
+            SearchLaneRequirementsCollector searchLaneRequirementsCollector
+        ) {
+            return createClusterInfoService(
+                new NoOpClient(threadPool),
+                estimatedHeapUsageCollector,
+                cacheSizesAndCommitmentCollector,
+                partitionSizeCollector,
+                searchLaneRequirementsCollector
+            );
+        }
+
+        InternalClusterInfoService createClusterInfoService(
+            NoOpClient client,
+            EstimatedHeapUsageCollector estimatedHeapUsageCollector,
+            CacheSizesAndCommitmentCollector cacheSizesAndCommitmentCollector,
+            PartitionSizeCollector partitionSizeCollector,
+            SearchLaneRequirementsCollector searchLaneRequirementsCollector
+        ) {
+            final InternalClusterInfoService clusterInfoService = new InternalClusterInfoService(
+                settings,
+                new WriteLoadConstraintSettings(clusterService.getClusterSettings()),
+                clusterService,
+                threadPool,
+                client,
+                estimatedHeapUsageCollector,
+                cacheSizesAndCommitmentCollector,
+                partitionSizeCollector,
+                NodeUsageStatsForThreadPoolsCollector.EMPTY,
+                searchLaneRequirementsCollector
+            );
+            // Refresh is a no-op if there are no listeners, and AsyncRefresh asserts that at least one listener is notified.
+            clusterInfoService.addListener(ignored -> {});
+            return clusterInfoService;
+        }
+
+        @Override
+        public void close() {
+            clusterService.close();
+        }
+    }
 }