KAFKA-20372: Fix flaky test in NamedTopologyIntegrationTest

streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/NamedTopologyIntegrationTest.java

@@ -657,7 +657,23 @@ public void shouldAddToEmptyInitialTopologyRemoveResetOffsetsThenAddSameNamedTop
 
         assertThat(waitUntilMinKeyValueRecordsReceived(consumerConfig, COUNT_OUTPUT, 5), equalTo(COUNT_OUTPUT_DATA));
         assertThat(waitUntilMinKeyValueRecordsReceived(consumerConfig, SUM_OUTPUT, 5), equalTo(SUM_OUTPUT_DATA));
+
+        TestUtils.waitForCondition(
+                () -> streams.allLocalTasksRunningForTopology(TOPOLOGY_1),
+                () -> "Not all local tasks for topology " + TOPOLOGY_1
+                        + " are initialized and in RUNNING state before remove. "
+                        + "streamsState=" + streams.state()
+                        + ", localThreads=" + streams.metadataForLocalThreads()
+        );
         streams.removeNamedTopology(TOPOLOGY_1, true).all().get();
+
+        TestUtils.waitForCondition(
+                () -> !streams.hasAnyLocalTaskForTopology(TOPOLOGY_1),
+                () -> "Topology " + TOPOLOGY_1
+                        + " still has local tasks after remove. "
+                        + "streamsState=" + streams.state()
+                        + ", localThreads=" + streams.metadataForLocalThreads()
+        );
         streams.cleanUpNamedTopology(TOPOLOGY_1);
 
         CLUSTER.getAllTopicsInCluster().stream().filter(t -> t.contains("-changelog") || t.contains("-repartition")).forEach(t -> {

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java

@@ -2078,6 +2078,17 @@ TaskManager taskManager() {
         return taskManager;
     }
 
+    // VisibleForTesting
+    public boolean hasAnyTaskForTopology(final String topologyName) {
+        return taskManager.hasAnyTaskForTopology(topologyName);
+    }
+
+
+    // VisibleForTesting
+    public boolean allTasksRunningForTopology(final String topologyName) {
+        return taskManager.allTasksRunningForTopology(topologyName);
+    }
+
     int currentNumIterations() {
         return numIterations;
     }

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

@@ -1631,6 +1631,38 @@ Map<TaskId, Task> allTasks() {
         return ret;
     }
 
+    // VisibleForTesting
+    boolean hasAnyTaskForTopology(final String topologyName) {
+        return allTasks().keySet().stream().anyMatch(taskId -> topologyName.equals(taskId.topologyName()));
+    }
+
+    /**
+     * Returns {@code true} if every task for the given topology is initialized and in
+     * {@link State#RUNNING}.
+     *
+     * <p>If there are no tasks for the given topology, this method returns {@code true}.
+     *
+     * @param topologyName the topology name
+     * @return {@code true} if all matching tasks are initialized and in {@link State#RUNNING},
+     *         or if there are no matching tasks; {@code false} otherwise
+     */
+    // VisibleForTesting
+    boolean allTasksRunningForTopology(final String topologyName) {
+        final Map<TaskId, Task> allTasks = allTasks();
+        final Set<TaskId> initializedTaskIds = tasks.allInitializedTaskIds();
+
+        for (final Map.Entry<TaskId, Task> entry : allTasks.entrySet()) {
+            final TaskId taskId = entry.getKey();
+            if (topologyName.equals(taskId.topologyName())) {
+                if (!initializedTaskIds.contains(taskId) || entry.getValue().state() != State.RUNNING) {
+                    return false;
+                }
+            }
+        }
+
+        return true;
+    }
+
     /**
      * Returns tasks owned by the stream thread.
      * This does not return any tasks currently owned by the state updater.

streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java

@@ -443,4 +443,18 @@ public Map<String, Map<Integer, LagInfo>> allLocalStorePartitionLagsForTopology(
                 .collect(Collectors.toList())));
         return allLocalStorePartitionLags(allTopologyTasks);
     }
+
+    // VisibleForTesting
+    public boolean hasAnyLocalTaskForTopology(final String topologyName) {
+        synchronized (threads) {
+            return threads.stream().anyMatch(thread -> thread.hasAnyTaskForTopology(topologyName));
+        }
+    }
+
+    // VisibleForTesting
+    public boolean allLocalTasksRunningForTopology(final String topologyName) {
+        synchronized (threads) {
+            return threads.stream().allMatch(thread -> thread.allTasksRunningForTopology(topologyName));
+        }
+    }
 }