Status

Current state: Under Discussion

Discussion thread: here and here

JIRA: here 

Please keep the discussion on the mailing list rather than commenting on the wiki (wiki discussions get unwieldy fast).

Motivation

KIP by Julien Brunet, Adam Souquières, Sébastien Viale, Marie-Laure Momplot

The current TopologyTestDriver only supports test input and output topics with a single partition, and therefore does not allow testing of topologies involving repartition operations.

The inability to test multi-partitioned streams leads to complex manual testing that cannot be automated or scripted — or, in many cases, results in missing tests altogether.

As a workaround, developers often rely on the EmbeddedKafkaCluster to run integration-style tests with multiple partitions. However, this approach requires spinning up a local Kafka cluster, managing configurations, and dealing with non-trivial setup and teardown logic. This makes it less user-friendly, slower, and harder to integrate into fast-running unit test suites.

For instance, when a key is modified before applying the .process() operator, Kafka Streams does not automatically create a repartition topic. This behavior can cause issues that go undetected when relying solely on single-partition unit tests.

This KIP proposes to introduce multi-partition support in the TopologyTestDriver, enabling more accurate and convenient stream testing while improving automated unit test coverage.

Public Interfaces

TopologyTestDriverBuilder (new class) 

TopologyTestDriverBuilder is the only supported way to create a TopologyTestDriver.

public class TopologyTestDriverBuilder {
    public TopologyTestDriverBuilder(Topology topology);
    public TopologyTestDriverBuilder withConfig(Properties config);
    public TopologyTestDriverBuilder withInitialWallClockTime(Instant initialWallClockTime);
    public TopologyTestDriverBuilder declareTopic(String topicName, int partitions);
    public TopologyTestDriver build();
}

TopologyTestDriver class changes 

  • Deprecate all public constructors 
@Deprecated
public TopologyTestDriver(final Topology topology)
@Deprecated
public TopologyTestDriver(final Topology topology, final Properties config)
@Deprecated
public TopologyTestDriver(final Topology topology, final Instant initialWallClockTimeMs)
@Deprecated
public TopologyTestDriver(final Topology topology, final Properties config, final Instant initialWallClockTimeMs)

They remain functional for backward compatibility but MUST NOT be used in new code and will be removed in a future major release.

  • Add new methods to get state stores by partition
// New per-partition store accessor ─────────────────────────────────
public StateStore getStateStore(final String name, final int partition) 
public <K, V> KeyValueStore<K, V> getKeyValueStore(final String name, final int partition)
public <K, V> KeyValueStore<K, ValueAndTimestamp<V>> getTimestampedKeyValueStore(final String name, final int partition)
public <K, V> VersionedKeyValueStore<K, V> getVersionedKeyValueStore(final String name, final int partition)
public <K, V> WindowStore<K, V> getWindowStore(final String name, final int partition)
public <K, V> WindowStore<K, ValueAndTimestamp<V>> getTimestampedWindowStore(final String name, final int partition)
public <K, V> SessionStore<K, V> getSessionStore(final String name, final int partition)

public <K, V> KeyValueStore<K, ValueTimestampHeaders<V>> getTimestampedKeyValueStoreWithHeaders(final String name, final int partition)
public <K, V> WindowStore<K, ValueTimestampHeaders<V>> getTimestampedWindowStoreWithHeaders(final String name, final int partition)
public <K, V> SessionStoreWithHeaders<K, V> getSessionStoreWithHeaders(final String name, final int partition)

TestRecord class changes 

  • Add a new constructor to TestRecord that accepts a partition number as an additional argument.

TestRecord(final K key, final V value, final Headers headers, final Instant recordTime, final int partition) {}
  • All existing constructors that do not take an explicit partition argument default partition to -1, indicating "no explicit partition set".

Update equals() and hashCode() to include the partition field.

- For input records, existing code that does not set a partition explicitly will produce partition = -1 on both sides of an equality check, preserving backward compatibility.

- For output records in single-partition mode, the partition field is set to -1 (not 0), to preserve backward compatibility with existing tests. Changing it to 0 would break all existing equality checks between output TestRecords and expected records built from partition-less constructors (which default to -1). Users who need to assert the partition explicitly in single-partition mode should use equalsIgnorePartition().

  • Add a utility method for partition-agnostic comparison:
public boolean equalsIgnorePartition(final TestRecord<K, V> o) {}

This allows tests that do not care about partition placement to use assertTrue(expected.equalsIgnorePartition(actual)).

Note that -1 never appears in output TestRecords , records read from readRecordsToList() always carry the real resolved partition.

-1 only exists as a sentinel value on input records where no explicit partition was set.

Proposed Changes

Internally, TopologyTestDriver will maintain topic metadata that includes:

  • The number of partitions for each topic

  • A task per partition to simulate the log structure

To avoid introducing unnecessary complexity when supporting multiple partitions and tasks, both in the TopologyTestDriver implementation and for users, we propose introducing a dedicated setup phase.

Instead of dynamically creating and updating tasks as topics are defined or as records are piped, all input and output topics must be created upfront, before any records can be processed.

The revised control flow would be:

Setup phase (builder chain):

TopologyTestDriverBuilder is the single entry point for all new tests.

The mode is determined by the declared topics: if at least one topic is declared with a partition count greater than 1, the driver operates in multi-partition mode; otherwise it operates in single-partition mode. 

  1.  Instantiate TopologyTestDriver with TopologyTestDriverBuilder.
  2. Declare all needed configurations via withConfig() and withInitialWallClockTime().
  3.  Declare topic partition counts via declareTopic(name, n).
  4.  Call build(), this creates all Tasks and GlobalTasks.

If no topics are declared, or all declared topics have a partition count of 1, the driver operates in single-partition mode, identical behavior to the current release.

An exception is thrown if a record is piped into an undeclared input topic while the driver is operating in multi-partition mode (that is, at least one topic has a partition count greater than 1).

Output topics do not need to be declared. Unlike input topics, they do not participate in task initialization and are therefore automatically created with a single partition.

Execution phase (identical in both modes): 

  1. Create input/output topic handles via createInputTopic / createOutputTopic. 
  2. Pipe records via pipeInput: the driver resolves the target partition  by explicit partition on the TestRecord, or by key hash (murmur2 % n,  consistent with BuiltInPartitioner).
  3. After each pipeInput, the driver drains all tasks to quiescence before  returning: tasks are processed in ascending stream-time order, with ascending (subtopologyId, partition) as a deterministic tie-breaker when stream-times are equal (e.g. immediately after a fan-out).
  4.  Read outputs via readRecordsToList(): filter by partition if needed.

Partition routing

When a record is piped via pipeInput, the target partition is resolved using the following priority order: 

  1. Explicit partition: if a partition number is set on the TestRecord, it is used directly. If it is out of range [0, n), an IllegalArgumentException is thrown.
  2. Null key: if no explicit partition is set and the key is null, the record is distributed via round-robin across all partitions of the topic.
  3. Key-based routing: if no explicit partition is set and the key is non-null, the partition is computed as:
    Utils.toPositive(Utils.murmur2(keyBytes)) % n which matches the behaviour of BuiltInPartitioner.

Cases 2 and 3 will be handled internally by a new TopologyTestDriverPartitioner that implements the Partitioner interface

Processing order on fan-out

When a single pipeInput causes records to fan out across multiple downstream tasks, the driver drains all tasks to quiescence before returning.

On each scheduling step, the task with the lowest current stream-time is selected next.

When multiple tasks share the same stream-time, which is the common case immediately after a fan-out, since all newly-enqueued downstream tasks start with no prior stream-time, the tie is broken by ascending (subtopologyId, partition), i.e. TaskId natural order.

This ordering is deterministic and stable, guaranteed by the TreeMap<TaskId, StreamTask> used internally to store tasks.

Store access in multi-partition mode

 In multi-partition mode, state stores are partitioned: each task owns its own store instance. 

  1. Partitioned stores (non-global): the no-argument accessors  getStateStore(), getKeyValueStore(), getSessionStore(), and getWindowStore() throw IllegalStateException in multi-partition mode,  since no single partition can be inferred. 
    The per-partition overload  must be used instead: 
       driver.getKeyValueStore("counts", 0); // partition 0
       driver.getKeyValueStore("counts", 1); // partition 1
  2. Global stores: global stores are not partitioned, they accumulate  records from all partitions into a single shared instance.
    The     no-argument accessors remain valid and unchanged in multi-partition  mode.

Summary of Contract

TopologyTestDriverBuilder is the recommended entry point for all new tests.
If no topics are declared, or if all topics have a partition count of 1, the driver operates in single-partition mode, identical behavior to the current release.
When one or more topics have multiple partitions, the driver creates one task per partition and enables partition-aware routing and store access.
Existing TopologyTestDriver constructors remain functional but are deprecated.

Example Usage Multi-Partition Mode

// Multi-partition mode (new)
TopologyTestDriver driver = new TopologyTestDriverBuilder(topology)
    .withConfig(config)
    .withInitialWallClockTime(Instant.now())
    .declareTopic("topic1", 3)
    .declareTopic("topic2", 3)
    .build();

TestInputTopic<String, String> input = driver.createInputTopic(
    "input", Serdes.String().serializer(), Serdes.String().serializer());
TestOutputTopic<String, String> output = driver.createOutputTopic(
    "output", Serdes.String().deserializer(), Serdes.String().deserializer());

// Pipe with explicit partition
input.pipeInput(new TestRecord<>("key0", "value0", 1000L, 0));
// Pipe with key-based routing (murmur2 % 3)
input.pipeInput("key1", "value1");
// Null key → distributed round-robin across partitions (partition = counter++ % 3)
// First call  → partition 0
// Second call → partition 1
// Third call  → partition 2
input.pipeInput(new TestRecord<>(null, "value2", 1000L));
input.pipeInput(new TestRecord<>(null, "value3", 1000L));
input.pipeInput(new TestRecord<>(null, "value4", 1000L));

// Read all outputs, filter by partition if needed
List<TestRecord<String, String>> all = output.readRecordsToList();
List<TestRecord<String, String>> p0  = all.stream()
    .filter(r -> r.getPartition() == 0)
    .collect(Collectors.toList());

Example Usage Single-Partition Mode

// Single-partition mode using TopologyTestDriverBuilder - no need to declare topics
TopologyTestDriver driver = new TopologyTestDriverBuilder(topology)
    .withConfig(config)
    .withInitialWallClockTime(Instant.now())
    .build();

TestInputTopic<String, String> input = driver.createInputTopic(
    "input", Serdes.String().serializer(), Serdes.String().serializer());
TestOutputTopic<String, String> output = driver.createOutputTopic(
    "output", Serdes.String().deserializer(), Serdes.String().deserializer());

input.pipeInput("key1", "value1"); // identical to current behaviour

Compatibility, Deprecation, and Migration Plan

  • Deprecation: All existing TopologyTestDriver constructors are deprecated in this release. They remain fully functional and no existing test will break. Users are encouraged to migrate to TopologyTestDriverBuilder.
// Before (deprecated)
new TopologyTestDriver(topology, config)

// After
new TopologyTestDriverBuilder(topology).withConfig(config).build()
  • Backward compatible: All existing TopologyTestDriver constructors, methods, and behaviors remain compatible. Existing tests require no modification. 
  • Single-partition mode preserves the behavior of the current release. If no topics are declared, or if all declared topics have a partition count of 1, the driver operates in single-partition mode and follows the existing execution path.
  • In single-partition mode, topics that are not declared via the builder implicitly use a single partition, preserving current semantics.
  • Multi-partition mode is enabled when at least one input topic is declared with a partition count greater than 1. In this mode, all input topics must be declared before build(). Attempting to pipe a record into an undeclared input topic results in an exception. Output topics do not need to be declared and default to a single partition.
  • All existing TestRecord constructors default partition to -1. Existing assertEquals calls between two records created without an explicit partition remain valid since both sides will carry partition = -1
  • Output TestRecords produced in single-partition mode carry partition -1 to preserve backward compatibility. Setting it to 0 would silently break all existing equality checks against expected records built from constructors that do not specify a partition (which default to -1)

Note:

In single-partition mode, the null-key rule is a no-op since partition 0 is the only partition. 

Test Plan

Add unit tests verifying:

  1. Multi-partition input topics route data correctly based on:

    • Explicit partition

    • Key-based partitioning

    • Null-key records are distributed via round-robin across all partitions, and never cause a NullPointerException.

  2. Repartition and join operations behave as expected across multiple partitions.

  3. Backward compatibility:  1 partition default still behaves identically to current implementation.

Rejected Alternatives

Defining the partition number directly in the TopologyTestDriver constructor would simplify the setup. However, this approach restricts testing to topics that share the same number of partitions. To better emulate the actual Kafka behavior, we opted to specify the partition number in the input and output topic declarations instead.

Adding a partition argument to TestInputTopic and TestOutputTopic would introduce many additional overloads. It would be cleaner and more consistent with the Producer API to specify the partition number on the TestRecord instead.

Dynamically updating the Tasks and GlobalTasks whenever an input topic is created with a partition number higher than the current maximum was also considered. However, this would introduce additional internal complexity and require mutating tasks at runtime, making the driver lifecycle less predictable.


  • No labels