vvcephei commented on a change in pull request #8588:
URL: https://github.com/apache/kafka/pull/8588#discussion_r424718290



##########
File path: 
streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java
##########
@@ -75,4 +94,303 @@
     static UUID uuidForInt(final int n) {
         return new UUID(0, n);
     }
+
+    static void assertValidAssignment(final int numStandbyReplicas,
+                                      final Set<TaskId> statefulTasks,
+                                      final Set<TaskId> statelessTasks,
+                                      final Map<UUID, ClientState> 
assignedStates,
+                                      final StringBuilder failureContext) {
+        assertValidAssignment(
+            numStandbyReplicas,
+            0,
+            statefulTasks,
+            statelessTasks,
+            assignedStates,
+            failureContext
+        );
+    }
+
+    static void assertValidAssignment(final int numStandbyReplicas,
+                                      final int maxWarmupReplicas,
+                                      final Set<TaskId> statefulTasks,
+                                      final Set<TaskId> statelessTasks,
+                                      final Map<UUID, ClientState> 
assignedStates,
+                                      final StringBuilder failureContext) {
+        final Map<TaskId, Set<UUID>> assignments = new TreeMap<>();
+        for (final TaskId taskId : statefulTasks) {
+            assignments.put(taskId, new TreeSet<>());
+        }
+        for (final TaskId taskId : statelessTasks) {
+            assignments.put(taskId, new TreeSet<>());
+        }
+        for (final Map.Entry<UUID, ClientState> entry : 
assignedStates.entrySet()) {
+            validateAndAddActiveAssignments(statefulTasks, statelessTasks, 
failureContext, assignments, entry);
+            validateAndAddStandbyAssignments(statefulTasks, statelessTasks, 
failureContext, assignments, entry);
+        }
+
+        final AtomicInteger remainingWarmups = new 
AtomicInteger(maxWarmupReplicas);
+
+        final TreeMap<TaskId, Set<UUID>> misassigned =
+            assignments
+                .entrySet()
+                .stream()
+                .filter(entry -> {
+                    final int expectedActives = 1;
+                    final boolean isStateless = 
statelessTasks.contains(entry.getKey());
+                    final int expectedStandbys = isStateless ? 0 : 
numStandbyReplicas;
+                    // We'll never assign even the expected number of standbys 
if they don't actually fit in the cluster
+                    final int expectedAssignments = Math.min(
+                        assignedStates.size(),
+                        expectedActives + expectedStandbys
+                    );
+                    final int actualAssignments = entry.getValue().size();
+                    if (actualAssignments == expectedAssignments) {
+                        return false; // not misassigned
+                    } else {
+                        if (actualAssignments == expectedAssignments + 1 && 
remainingWarmups.get() > 0) {
+                            remainingWarmups.getAndDecrement();
+                            return false; // it's a warmup, so it's fine
+                        } else {
+                            return true; // misassigned
+                        }
+                    }
+                })
+                .collect(entriesToMap(TreeMap::new));
+
+        if (!misassigned.isEmpty()) {

Review comment:
       L131-158 is just gathering the information about whether each task is 
correctly assigned or not, based on its type and the standby configs (and maybe 
the warmup config). It doesn't make any assertions. So this check is actually 
the assertion, that no tasks are incorrectly assigned.
   
   Doing it this way is nicer, since when it fails, it tells you _all_ the 
incorrectly assigned tasks, not just the first one.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to