Wait longer to reduce flakiness of TestNodeStateManager

Author: losipiuk

core/trino-main/src/test/java/io/trino/server/TestNodeStateManager.java

@@ -60,7 +60,7 @@
 
 class TestNodeStateManager
 {
-    public static final int GRACE_PERIOD_MILLIS = 200;
+    public static final int GRACE_PERIOD_MILLIS = 1000;
     private FakeScheduledExecutorService executor;
     private NodeStateManager nodeStateManager;
     private TestingTicker ticker;
@@ -94,7 +94,7 @@ void testDrain()
         ticker.increment(1, SECONDS);
         executor.run();
 
-        await().atMost(1, SECONDS).untilAsserted(() -> assertThat(nodeStateManager.getServerState()).isEqualTo(DRAINED));
+        await().atMost(5, SECONDS).untilAsserted(() -> assertThat(nodeStateManager.getServerState()).isEqualTo(DRAINED));
     }
 
     @Test
@@ -105,8 +105,8 @@ void testTransitionToShuttingDown()
         nodeStateManager.transitionState(NodeState.SHUTTING_DOWN);
         assertThat(nodeStateManager.getServerState()).isEqualTo(NodeState.SHUTTING_DOWN);
 
-        // here wait for at least 2 grace periods, and add some slack to reduce test flakyness
-        await().atMost(4 * GRACE_PERIOD_MILLIS + 100, MILLISECONDS).until(() -> shutdownAction.isShuttingDown());
+        // here wait for at least 4 grace periods, and add some slack to reduce test flakyness
+        await().atMost(4 * GRACE_PERIOD_MILLIS + 1000, MILLISECONDS).until(() -> shutdownAction.isShuttingDown());
     }
 
     @Test
@@ -117,7 +117,7 @@ void testCannotReactivateShuttingDown()
         nodeStateManager.transitionState(NodeState.SHUTTING_DOWN);
         assertThat(nodeStateManager.getServerState()).isEqualTo(NodeState.SHUTTING_DOWN);
 
-        // here wait for at least 2 grace periods, and add some slack to reduce test flakyness
+        // here wait for at least 4 grace periods, and add some slack to reduce test flakyness
         await().atMost(4 * GRACE_PERIOD_MILLIS, MILLISECONDS).until(() -> shutdownAction.isShuttingDown());
 
         assertThatThrownBy(() -> nodeStateManager.transitionState(ACTIVE))
@@ -144,9 +144,8 @@ void testImmediateTransitionToShuttingDownWhenDrained()
         nodeStateManager.transitionState(NodeState.SHUTTING_DOWN);
         assertThat(nodeStateManager.getServerState()).isEqualTo(NodeState.SHUTTING_DOWN);
 
-        // here only wait for minimal amount of time, as shutdown should be immediate
         await().pollInterval(1, MILLISECONDS)
-                .atMost(100, MILLISECONDS).until(() -> shutdownAction.isShuttingDown());
+                .atMost(1, SECONDS).until(() -> shutdownAction.isShuttingDown());
     }
 
     @Test
@@ -171,7 +170,7 @@ void testWaitActiveTasksToFinishDuringShutdown()
         // make sure that nodeStateManager registered a listener for tasks to finish
         ticker.increment(1, SECONDS);
         executor.run();
-        await().atMost(1, SECONDS).until(() -> sqlTasksObservable.getTasks().size() == 1);
+        await().atMost(5, SECONDS).until(() -> sqlTasksObservable.getTasks().size() == 1);
 
         // simulate task completion after some time
         tasks.set(Collections.emptyList());
@@ -213,7 +212,7 @@ void testWaitActiveTasksToFinishDuringDraining()
                 .stateChanged(TaskState.FINISHED);
 
         // when NodeStateManager sees task finished - it will drain after another drain period
-        await().atMost(1, SECONDS)
+        await().atMost(5, SECONDS)
                 .untilAsserted(() -> assertThat(nodeStateManager.getServerState()).isEqualTo(DRAINED));
     }