Fix testReturnsOnlyRequestedStats after merging 124902

Author: JeremyDahlgren

server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/TransportGetAllocationStatsActionTests.java

@@ -43,22 +43,20 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Consumer;
 
 import static org.hamcrest.Matchers.anEmptyMap;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.not;
 import static org.mockito.Mockito.mock;
-
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class TransportGetAllocationStatsActionTests extends ESTestCase {
 
     private long startTimeMillis;
-    private TimeValue cacheTTL;
+    private TimeValue allocationStatsCacheTTL;
     private ControlledRelativeTimeThreadPool threadPool;
     private ClusterService clusterService;
     private TransportService transportService;
@@ -71,12 +69,14 @@ public class TransportGetAllocationStatsActionTests extends ESTestCase {
     public void setUp() throws Exception {
         super.setUp();
         startTimeMillis = 0L;
-        cacheTTL = TimeValue.timeValueMinutes(1);
+        allocationStatsCacheTTL = TimeValue.timeValueMinutes(1);
         threadPool = new ControlledRelativeTimeThreadPool(TransportClusterAllocationExplainActionTests.class.getName(), startTimeMillis);
         clusterService = ClusterServiceUtils.createClusterService(
             threadPool,
             new ClusterSettings(
-                Settings.builder().put(TransportGetAllocationStatsAction.CACHE_TTL_SETTING.getKey(), cacheTTL.toString()).build(),
+                Settings.builder()
+                    .put(TransportGetAllocationStatsAction.CACHE_TTL_SETTING.getKey(), allocationStatsCacheTTL.toString())
+                    .build(),
                 ClusterSettings.BUILT_IN_CLUSTER_SETTINGS
             )
         );
@@ -107,7 +107,17 @@ public void tearDown() throws Exception {
         transportService.close();
     }
 
+    private void disableAllocationStatsCache() {
+        setAllocationStatsCacheTTL(TimeValue.ZERO);
+    }
+
+    private void setAllocationStatsCacheTTL(TimeValue ttl) {
+        clusterService.getClusterSettings()
+            .applySettings(Settings.builder().put(TransportGetAllocationStatsAction.CACHE_TTL_SETTING.getKey(), ttl.toString()).build());
+    };
+
     public void testReturnsOnlyRequestedStats() throws Exception {
+        disableAllocationStatsCache();
         int expectedNumberOfStatsServiceCalls = 0;
 
         for (final var metrics : List.of(
@@ -149,6 +159,7 @@ public void testReturnsOnlyRequestedStats() throws Exception {
     }
 
     public void testDeduplicatesStatsComputations() throws InterruptedException {
+        disableAllocationStatsCache();
         final var requestCounter = new AtomicInteger();
         final var isExecuting = new AtomicBoolean();
         when(allocationStatsService.stats()).thenAnswer(invocation -> {
@@ -205,13 +216,6 @@ public void testGetStatsWithCachingEnabled() throws Exception {
             when(allocationStatsService.stats()).thenReturn(stats);
         };
 
-        final Consumer<TimeValue> setCacheTTL = (ttl) -> {
-            clusterService.getClusterSettings()
-                .applySettings(
-                    Settings.builder().put(TransportGetAllocationStatsAction.CACHE_TTL_SETTING.getKey(), ttl.toString()).build()
-                );
-        };
-
         final CheckedConsumer<ActionListener<Void>, Exception> threadTask = l -> {
             final var request = new TransportGetAllocationStatsAction.Request(
                 TEST_REQUEST_TIMEOUT,
@@ -233,22 +237,22 @@ public void testGetStatsWithCachingEnabled() throws Exception {
         verify(allocationStatsService, times(++numExpectedAllocationStatsServiceCalls)).stats();
 
         // Force the cached stats to expire.
-        threadPool.setCurrentTimeInMillis(startTimeMillis + (2 * cacheTTL.getMillis()));
+        threadPool.setCurrentTimeInMillis(startTimeMillis + (2 * allocationStatsCacheTTL.getMillis()));
 
         // Expect a single call to the stats service on the cache miss.
         resetExpectedAllocationStats.run();
         ESTestCase.startInParallel(between(1, 5), threadNumber -> safeAwait(threadTask));
         verify(allocationStatsService, times(++numExpectedAllocationStatsServiceCalls)).stats();
 
         // Update the TTL setting to disable the cache, we expect a service call each time.
-        setCacheTTL.accept(TimeValue.ZERO);
+        setAllocationStatsCacheTTL(TimeValue.ZERO);
         threadTask.accept(ActionListener.noop());
         threadTask.accept(ActionListener.noop());
         numExpectedAllocationStatsServiceCalls += 2;
         verify(allocationStatsService, times(numExpectedAllocationStatsServiceCalls)).stats();
 
         // Re-enable the cache, only one thread should call the stats service.
-        setCacheTTL.accept(TimeValue.timeValueMinutes(5));
+        setAllocationStatsCacheTTL(TimeValue.timeValueMinutes(5));
         resetExpectedAllocationStats.run();
         ESTestCase.startInParallel(between(1, 5), threadNumber -> safeAwait(threadTask));
         verify(allocationStatsService, times(++numExpectedAllocationStatsServiceCalls)).stats();