[GSoC 2026] Kafka Streams runner — KafkaStreamsTestRunner test harness #39211
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,201 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package org.apache.beam.runners.kafka.streams; | ||
|
|
||
| import java.time.Duration; | ||
| import java.util.HashSet; | ||
| import java.util.List; | ||
| import java.util.Properties; | ||
| import java.util.Set; | ||
| import java.util.UUID; | ||
| import org.apache.beam.model.pipeline.v1.RunnerApi; | ||
| import org.apache.beam.runners.fnexecution.provisioning.JobInfo; | ||
| import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator; | ||
| import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext; | ||
| import org.apache.beam.sdk.Pipeline; | ||
| import org.apache.beam.sdk.options.PipelineOptions; | ||
| import org.apache.beam.sdk.options.PipelineOptionsFactory; | ||
| import org.apache.beam.sdk.options.PortablePipelineOptions; | ||
| import org.apache.beam.sdk.testing.CrashingRunner; | ||
| import org.apache.beam.sdk.util.construction.Environments; | ||
| import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; | ||
| import org.apache.beam.sdk.util.construction.PipelineTranslation; | ||
| import org.apache.kafka.common.serialization.ByteArrayDeserializer; | ||
| import org.apache.kafka.common.serialization.ByteArraySerializer; | ||
| import org.apache.kafka.common.serialization.Serdes; | ||
| import org.apache.kafka.streams.StreamsConfig; | ||
| import org.apache.kafka.streams.TestInputTopic; | ||
| import org.apache.kafka.streams.TestOutputTopic; | ||
| import org.apache.kafka.streams.Topology; | ||
| import org.apache.kafka.streams.TopologyDescription; | ||
| import org.apache.kafka.streams.TopologyTestDriver; | ||
| import org.apache.kafka.streams.test.TestRecord; | ||
|
|
||
| /** | ||
| * Test harness that runs a Beam {@link Pipeline} through the Kafka Streams runner's translation and | ||
| * a {@link TopologyTestDriver}, so tests do not repeat the translate + drive boilerplate. | ||
| * | ||
| * <p>Usage: build a pipeline with {@link #testOptions()}, then call {@link #run(Pipeline)}. Side | ||
| * effects (e.g. a {@code SharedTestCollector} written by a recording DoFn) have completed when it | ||
| * returns. | ||
| * | ||
| * <p>{@link TopologyTestDriver} does not loop a low-level sink topic back into its source, so an | ||
| * internal repartition topic (one that is both a sink and a source in the topology — e.g. the one | ||
| * GroupByKey introduces) would otherwise dead-end. {@link #run(Pipeline)} discovers those topics | ||
| * from the {@link TopologyDescription} and round-trips them until no more records flow, standing in | ||
| * for the broker. | ||
| */ | ||
| public final class KafkaStreamsTestRunner { | ||
|
|
||
| private static final int MAX_ROUND_TRIPS = 100; | ||
|
|
||
| private KafkaStreamsTestRunner() {} | ||
|
|
||
| /** Pipeline options for a Kafka Streams runner test: the EMBEDDED harness and a unique app id. */ | ||
| public static PipelineOptions testOptions() { | ||
| String applicationId = "ks-test-" + UUID.randomUUID(); | ||
| PipelineOptions options = | ||
| PipelineOptionsFactory.fromArgs("--applicationId=" + applicationId).create(); | ||
| options.setRunner(CrashingRunner.class); | ||
| options.as(KafkaStreamsPipelineOptions.class).setApplicationId(applicationId); | ||
| options | ||
| .as(PortablePipelineOptions.class) | ||
| .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); | ||
| return options; | ||
| } | ||
|
|
||
| /** | ||
| * Translates the pipeline into a Kafka Streams {@link KafkaStreamsTranslationContext}. Tests that | ||
| * need the {@link Topology} (e.g. to attach a capture processor before driving) use this and | ||
| * build their own {@link TopologyTestDriver}; simpler tests use {@link #run(Pipeline)}. | ||
| */ | ||
| public static KafkaStreamsTranslationContext translate(Pipeline pipeline) { | ||
| KafkaStreamsPipelineOptions options = | ||
| pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); | ||
| RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); | ||
| KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); | ||
| JobInfo jobInfo = | ||
| JobInfo.create( | ||
| options.getApplicationId(), | ||
| options.getJobName(), | ||
| "", | ||
| PipelineOptionsTranslation.toProto(options)); | ||
| KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); | ||
| translator.translate(context, translator.prepareForTranslation(pipelineProto)); | ||
| return context; | ||
| } | ||
|
|
||
| /** Translates and drives the pipeline to quiescence through a {@link TopologyTestDriver}. */ | ||
| public static void run(Pipeline pipeline) { | ||
| KafkaStreamsTranslationContext context = translate(pipeline); | ||
| Topology topology = context.getTopology(); | ||
| try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig(pipeline))) { | ||
| // Fire the Impulse wall-clock punctuator and let the initial records flow. | ||
| driver.advanceWallClockTime(Duration.ofSeconds(1)); | ||
| driver.advanceWallClockTime(Duration.ofSeconds(1)); | ||
| roundTripInternalTopics(driver, internalTopics(topology)); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * The name of some processor node with no successors (a topology leaf). Tests with a single leaf | ||
| * attach a capture processor here to observe what the last stage forwards; if a topology has more | ||
| * than one leaf, which one is returned is unspecified. | ||
| */ | ||
| public static String findAnyLeafProcessorName(Topology topology) { | ||
| for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { | ||
| for (TopologyDescription.Node node : subtopology.nodes()) { | ||
| if (node instanceof TopologyDescription.Processor && node.successors().isEmpty()) { | ||
| return node.name(); | ||
| } | ||
| } | ||
| } | ||
| throw new IllegalStateException("no leaf processor found in topology"); | ||
| } | ||
|
|
||
| /** Repartition/internal topics are the ones that appear as both a sink and a source. */ | ||
| private static Set<String> internalTopics(Topology topology) { | ||
| Set<String> sinkTopics = new HashSet<>(); | ||
| Set<String> sourceTopics = new HashSet<>(); | ||
| for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { | ||
| for (TopologyDescription.Node node : subtopology.nodes()) { | ||
| if (node instanceof TopologyDescription.Sink) { | ||
| String topic = ((TopologyDescription.Sink) node).topic(); | ||
| if (topic != null) { | ||
| sinkTopics.add(topic); | ||
| } | ||
| } else if (node instanceof TopologyDescription.Source) { | ||
| sourceTopics.addAll(((TopologyDescription.Source) node).topicSet()); | ||
| } | ||
| } | ||
| } | ||
| sinkTopics.retainAll(sourceTopics); | ||
| return sinkTopics; | ||
| } | ||
|
|
||
| /** | ||
| * Simulates the broker for internal repartition topics. | ||
| * | ||
| * <p>The runner shuffles data (and the watermark) through internal topics that a processor both | ||
| * writes to (a sink) and reads back from (a source) — e.g. the topic GroupByKey introduces to | ||
| * partition by key. On a real broker those records make the round trip automatically, but {@link | ||
| * TopologyTestDriver} does not connect a sink back to a source, so the downstream half of the | ||
| * topology would never see them. This drains what each internal topic's sink wrote and pipes it | ||
| * into that topic's source, repeating until nothing new flows (a fixpoint), which stands in for | ||
| * the broker and lets the pipeline run to completion. | ||
| */ | ||
| private static void roundTripInternalTopics(TopologyTestDriver driver, Set<String> topics) { | ||
|
Contributor
There was a problem hiding this comment. The purpose of this method is not clear to me, why do we feed outputs back to inputs?
Contributor
Author
There was a problem hiding this comment.
|
||
| for (int round = 0; round < MAX_ROUND_TRIPS; round++) { | ||
| boolean progressed = false; | ||
| for (String topic : topics) { | ||
| TestOutputTopic<byte[], byte[]> out = | ||
| driver.createOutputTopic( | ||
| topic, new ByteArrayDeserializer(), new ByteArrayDeserializer()); | ||
| List<TestRecord<byte[], byte[]>> records = out.readRecordsToList(); | ||
| if (records.isEmpty()) { | ||
| continue; | ||
| } | ||
| progressed = true; | ||
| TestInputTopic<byte[], byte[]> in = | ||
| driver.createInputTopic(topic, new ByteArraySerializer(), new ByteArraySerializer()); | ||
| for (TestRecord<byte[], byte[]> record : records) { | ||
| in.pipeInput(record); | ||
| } | ||
| } | ||
| if (!progressed) { | ||
| return; | ||
| } | ||
| } | ||
| throw new IllegalStateException( | ||
| "Internal topics did not reach quiescence after " + MAX_ROUND_TRIPS + " round trips"); | ||
| } | ||
|
Comment on lines
+162
to
+186
Contributor
There was a problem hiding this comment. Creating private static void roundTripInternalTopics(TopologyTestDriver driver, Set<String> topics) {
java.util.Map<String, TestOutputTopic<byte[], byte[]>> outputs = new java.util.HashMap<>();
java.util.Map<String, TestInputTopic<byte[], byte[]>> inputs = new java.util.HashMap<>();
for (String topic : topics) {
outputs.put(
topic,
driver.createOutputTopic(
topic, new ByteArrayDeserializer(), new ByteArrayDeserializer()));
inputs.put(
topic,
driver.createInputTopic(topic, new ByteArraySerializer(), new ByteArraySerializer()));
}
for (int round = 0; round < MAX_ROUND_TRIPS; round++) {
boolean progressed = false;
for (String topic : topics) {
TestOutputTopic<byte[], byte[]> out = outputs.get(topic);
List<TestRecord<byte[], byte[]>> records = out.readRecordsToList();
if (records.isEmpty()) {
continue;
}
progressed = true;
TestInputTopic<byte[], byte[]> in = inputs.get(topic);
for (TestRecord<byte[], byte[]> record : records) {
in.pipeInput(record);
}
}
if (!progressed) {
return;
}
}
throw new IllegalStateException(
"Internal topics did not reach quiescence after " + MAX_ROUND_TRIPS + " round trips");
} |
||
|
|
||
| /** Kafka Streams config for a {@link TopologyTestDriver} built from the pipeline's app id. */ | ||
| public static Properties streamsConfig(Pipeline pipeline) { | ||
| String applicationId = | ||
| pipeline.getOptions().as(KafkaStreamsPipelineOptions.class).getApplicationId(); | ||
| Properties props = new Properties(); | ||
| props.put(StreamsConfig.APPLICATION_ID_CONFIG, applicationId); | ||
| props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); | ||
| props.put( | ||
| StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); | ||
| props.put( | ||
| StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); | ||
| return props; | ||
| } | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If a source node in the topology is configured using a
Patterninstead of a set of topic names,topicSet()will returnnull. CallingsourceTopics.addAll(null)will throw aNullPointerException. Adding a null check before callingaddAllprevents potential crashes when pattern-based sources are used.