gemini-code-assist[bot] commented on code in PR #39211: URL: https://github.com/apache/beam/pull/39211#discussion_r3519118520
########## runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java: ########## @@ -0,0 +1,190 @@ +/* + * 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 the single processor node with no successors (the topology leaf). Tests attach a + * capture processor here to observe what the last stage forwards. + */ + public static String leafProcessorName(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"); + } Review Comment:  The iteration order of `subtopology.nodes()` (which returns a `Set`) is not guaranteed to be deterministic. If a topology contains multiple leaf processors (processors with no successors), `leafProcessorName` will return whichever leaf it encounters first, which could lead to non-deterministic test failures. Collecting all leaf processors and asserting that exactly one exists makes the test harness more robust and deterministic. ```java public static String leafProcessorName(Topology topology) { List<String> leaves = new java.util.ArrayList<>(); for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { for (TopologyDescription.Node node : subtopology.nodes()) { if (node instanceof TopologyDescription.Processor && node.successors().isEmpty()) { leaves.add(node.name()); } } } if (leaves.isEmpty()) { throw new IllegalStateException("no leaf processor found in topology"); } if (leaves.size() > 1) { throw new IllegalStateException("multiple leaf processors found in topology: " + leaves); } return leaves.get(0); } ``` ########## runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java: ########## @@ -0,0 +1,190 @@ +/* + * 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 the single processor node with no successors (the topology leaf). Tests attach a + * capture processor here to observe what the last stage forwards. + */ + public static String leafProcessorName(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()); + } Review Comment:  If a source node in the topology is configured using a `Pattern` instead of a set of topic names, `topicSet()` will return `null`. Calling `sourceTopics.addAll(null)` will throw a `NullPointerException`. Adding a null check before calling `addAll` prevents potential crashes when pattern-based sources are used. ```java } else if (node instanceof TopologyDescription.Source) { Set<String> topics = ((TopologyDescription.Source) node).topicSet(); if (topics != null) { sourceTopics.addAll(topics); } } ``` ########## runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java: ########## @@ -0,0 +1,190 @@ +/* + * 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 the single processor node with no successors (the topology leaf). Tests attach a + * capture processor here to observe what the last stage forwards. + */ + public static String leafProcessorName(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; + } + + /** Drains each internal topic's output and feeds it back into the source until nothing flows. */ + private static void roundTripInternalTopics(TopologyTestDriver driver, Set<String> topics) { + 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"); + } Review Comment:  Creating `TestInputTopic` and `TestOutputTopic` instances inside the round-trip loop on every iteration is inefficient as it repeatedly allocates new objects and performs lookups in `TopologyTestDriver`'s internal maps. Instantiating them once outside the loop improves performance and reduces overhead. ```java 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"); } ``` -- 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. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
