je-ik commented on code in PR #39211: URL: https://github.com/apache/beam/pull/39211#discussion_r3524754179
########## 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: This might be OK, but I'd suggest either rename it to "findAnyLeafProcessorName" or we could return List of nodes and caller might choose what to do. ########## 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) { Review Comment: The purpose of this method is not clear to me, why do we feed outputs back to inputs? -- 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]
