KAFKA-6145: KIP-441: Add test scenarios to ensure rebalance convergence

clients/src/main/java/org/apache/kafka/common/utils/Utils.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.utils;
 
+import java.util.EnumSet;
 import java.util.SortedSet;
 import java.util.TreeSet;
 import org.apache.kafka.common.KafkaException;
@@ -60,9 +61,13 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
+import java.util.function.BinaryOperator;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collector;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -1093,4 +1098,52 @@ public static <K1, V1, K2, V2> Map<K2, V2> transformMap(
             )
         );
     }
+
+    /**
+     * A Collector that offers two kinds of convenience:
+     * 1. You can specify the concrete type of the returned Map
+     * 2. You can turn a stream of Entries directly into a Map without having to mess with a key function
+     *    and a value function. In particular, this is handy if all you need to do is apply a filter to a Map's entries.
+     *
+     *
+     * One thing to be wary of: These types are too "distant" for IDE type checkers to warn you if you
+     * try to do something like build a TreeMap of non-Comparable elements. You'd get a runtime exception for that.
+     *
+     * @param mapSupplier The constructor for your concrete map type.
+     * @param <K> The Map key type
+     * @param <V> The Map value type
+     * @param <M> The type of the Map itself.
+     * @return new Collector<Map.Entry<K, V>, M, M>
+     */
+    public static <K, V, M extends Map<K, V>> Collector<Map.Entry<K, V>, M, M> entriesToMap(final Supplier<M> mapSupplier) {
+        return new Collector<Map.Entry<K, V>, M, M>() {
+            @Override
+            public Supplier<M> supplier() {
+                return mapSupplier;
+            }
+
+            @Override
+            public BiConsumer<M, Map.Entry<K, V>> accumulator() {
+                return (map, entry) -> map.put(entry.getKey(), entry.getValue());
+            }
+
+            @Override
+            public BinaryOperator<M> combiner() {
+                return (map, map2) -> {
+                    map.putAll(map2);
+                    return map;
+                };
+            }
+
+            @Override
+            public Function<M, M> finisher() {
+                return map -> map;
+            }
+
+            @Override
+            public Set<Characteristics> characteristics() {
+                return EnumSet.of(Characteristics.UNORDERED, Characteristics.IDENTITY_FINISH);
+            }
+        };
+    }
 }

streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java

@@ -63,24 +63,6 @@ public class AssignmentTestUtils {
      * @return the UUID created by repeating the digit n in the UUID format
      */
     static UUID uuidForInt(final Integer n) {
-        if (n < 1 || n > 7) {
-            throw new IllegalArgumentException("Must pass in a single digit number to the uuid builder, got n = " + n);
-        }
-        final StringBuilder builder = new StringBuilder(36);
-        for (int i = 0; i < 8; ++i) {
-            builder.append(n);
-        }
-        builder.append('-');
-
-        for (int i = 0; i < 3; ++i) {
-            for (int j = 0; j < 4; ++j) {
-                builder.append(n);
-            }
-            builder.append('-');
-        }
-        for (int i = 0; i < 12; ++i) {
-            builder.append(n);
-        }
-        return UUID.fromString(builder.toString());
+        return new UUID(0, n);
     }
 }