MINOR: Fix flakiness in state updater unit tests (#16562)

Unit test shouldRestoreSingleActiveStatefulTask() in DefaultStateUpdaterTest is flaky.

The flakiness comes from the fact that the state updater thread could call the first time changelogReader.allChangelogsCompleted() before it calls the first time changelogReader.completedChangelogs(). That happens, if runOnce() is run before the state updater thread reads a task from the input queue.

This commit fixes the flakiness, by making stubbing changelogReader.allChangelogsCompleted() depend on stubbing changelogReader.completedChangelogs().

Reviewers: Lucas Brutschy <[email protected]>, Matthias J. Sax <[email protected]>

streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStateUpdaterTest.java

@@ -109,7 +109,7 @@ class DefaultStateUpdaterTest {
     private final StreamsConfig config = new StreamsConfig(configProps(COMMIT_INTERVAL));
     private final ChangelogReader changelogReader = mock(ChangelogReader.class);
     private final TopologyMetadata topologyMetadata = unnamedTopology().build();
-    private DefaultStateUpdater stateUpdater =
+    private final DefaultStateUpdater stateUpdater =
         new DefaultStateUpdater("test-state-updater", metrics, config, null, changelogReader, topologyMetadata, time);
 
     @AfterEach
@@ -324,14 +324,15 @@ private void shouldImmediatelyAddStatelessTasksToRestoredTasks(final StreamTask.
     public void shouldRestoreSingleActiveStatefulTask() throws Exception {
         final StreamTask task =
             statefulTask(TASK_0_0, Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0)).inState(State.RESTORING).build();
+        final AtomicBoolean allChangelogCompleted = new AtomicBoolean(false);
         when(changelogReader.completedChangelogs())
             .thenReturn(Collections.emptySet())
             .thenReturn(Set.of(TOPIC_PARTITION_A_0))
-            .thenReturn(Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0));
-        when(changelogReader.allChangelogsCompleted())
-            .thenReturn(false)
-            .thenReturn(false)
-            .thenReturn(true);
+            .thenAnswer(invocation -> {
+                allChangelogCompleted.set(true);
+                return Set.of(TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0);
+            });
+        when(changelogReader.allChangelogsCompleted()).thenAnswer(invocation -> allChangelogCompleted.get());
         stateUpdater.start();
 
         stateUpdater.add(task);
@@ -362,7 +363,7 @@ public void shouldRestoreMultipleActiveStatefulTasks() throws Exception {
                 allChangelogCompleted.set(true);
                 return Set.of(TOPIC_PARTITION_C_0, TOPIC_PARTITION_A_0, TOPIC_PARTITION_B_0);
             });
-        when(changelogReader.allChangelogsCompleted()).thenReturn(allChangelogCompleted.get());
+        when(changelogReader.allChangelogsCompleted()).thenAnswer(invocation -> allChangelogCompleted.get());
         stateUpdater.start();
 
         stateUpdater.add(task1);