From 428dfa58a9b56eeb48634ccd01f8ed9a57a03ce5 Mon Sep 17 00:00:00 2001 From: Thomas Groh Date: Tue, 31 May 2016 17:09:19 -0700 Subject: [PATCH 1/3] Remove DirectRunner implementations from DoFnTester Part of the removal of runners-core from the Core SDK, and the deletion of the DirectRunner from the core SDK. --- .../beam/sdk/transforms/DoFnTester.java | 282 +++++++++++++++--- 1 file changed, 240 insertions(+), 42 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java index 16dc731816f8..569b86b87157 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java @@ -20,22 +20,20 @@ import org.apache.beam.sdk.annotations.Experimental; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.util.DirectModeExecutionContext; -import org.apache.beam.sdk.util.DirectSideInputReader; -import org.apache.beam.sdk.util.DoFnRunner; -import org.apache.beam.sdk.util.DoFnRunnerBase; -import org.apache.beam.sdk.util.DoFnRunners; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.util.PTuple; import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.sdk.util.UserCodeException; import org.apache.beam.sdk.util.WindowedValue; -import org.apache.beam.sdk.util.WindowingStrategy; -import org.apache.beam.sdk.util.common.Counter; -import org.apache.beam.sdk.util.common.CounterSet; +import org.apache.beam.sdk.util.WindowingInternals; import org.apache.beam.sdk.values.PCollectionView; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import com.google.common.base.Function; +import com.google.common.base.MoreObjects; import com.google.common.base.Objects; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; @@ -44,6 +42,7 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -198,7 +197,13 @@ public final List processBatch(InputT... inputElements) { public void startBundle() { resetState(); initializeState(); - fnRunner.startBundle(); + try { + TestContext context = createContext(fn); + context.setupDelegateAggregators(); + fn.startBundle(context); + } catch (Exception e) { + throw UserCodeException.wrap(e); + } state = State.STARTED; } @@ -220,7 +225,11 @@ public void processElement(InputT element) { if (state == State.UNSTARTED) { startBundle(); } - fnRunner.processElement(WindowedValue.valueInGlobalWindow(element)); + try { + fn.processElement(createProcessContext(fn, element)); + } catch (Exception e) { + throw UserCodeException.wrap(e); + } } /** @@ -239,7 +248,11 @@ public void finishBundle() { if (state == State.UNSTARTED) { startBundle(); } - fnRunner.finishBundle(); + try { + fn.finishBundle(createContext(fn)); + } catch (Exception e) { + throw UserCodeException.wrap(e); + } state = State.FINISHED; } @@ -276,8 +289,7 @@ public OutputT apply(OutputElementWithTimestamp input) { @Experimental public List> peekOutputElementsWithTimestamp() { // TODO: Should we return an unmodifiable list? - return Lists.transform( - outputManager.getOutput(mainOutputTag), + return Lists.transform(getOutput(mainOutputTag), new Function>() { @Override @SuppressWarnings("unchecked") @@ -336,8 +348,7 @@ public List> takeOutputElementsWithTimestamp */ public List peekSideOutputElements(TupleTag tag) { // TODO: Should we return an unmodifiable list? - return Lists.transform( - outputManager.getOutput(tag), + return Lists.transform(getOutput(tag), new Function, T>() { @SuppressWarnings("unchecked") @Override @@ -372,11 +383,15 @@ public List takeSideOutputElements(TupleTag tag) { * Returns the value of the provided {@link Aggregator}. */ public AggregateT getAggregatorValue(Aggregator agg) { + return extractAggregatorValue(agg, agg.getCombineFn()); + } + + private AggregateT extractAggregatorValue( + Aggregator agg, + CombineFn combiner) { @SuppressWarnings("unchecked") - Counter counter = - (Counter) - counterSet.getExistingCounter("user-" + STEP_NAME + "-" + agg.getName()); - return counter.getAggregate(); + AccumT accumulator = (AccumT) accumulators.get(agg.getName()); + return combiner.extractOutput(accumulator); } /** @@ -415,6 +430,205 @@ public int hashCode() { } } + private List> getOutput(TupleTag tag) { + @SuppressWarnings({"unchecked", "rawtypes"}) + List> elems = (List) outputs.get(tag); + return MoreObjects.firstNonNull(elems, Collections.>emptyList()); + } + + private TestContext createContext(DoFn fn) { + return new TestContext<>(fn, options, mainOutputTag, outputs, accumulators); + } + + private static class TestContext extends DoFn.Context { + private final PipelineOptions opts; + private final TupleTag mainOutputTag; + private final Map, List>> outputs; + private final Map accumulators; + + public TestContext( + DoFn fn, + PipelineOptions opts, + TupleTag mainOutputTag, + Map, List>> outputs, + Map accumulators) { + fn.super(); + this.opts = opts; + this.mainOutputTag = mainOutputTag; + this.outputs = outputs; + this.accumulators = accumulators; + } + + @Override + public PipelineOptions getPipelineOptions() { + return opts; + } + + @Override + public void output(OutT output) { + sideOutput(mainOutputTag, output); + } + + @Override + public void outputWithTimestamp(OutT output, Instant timestamp) { + sideOutputWithTimestamp(mainOutputTag, output, timestamp); + } + + @Override + protected Aggregator createAggregatorInternal( + final String name, final CombineFn combiner) { + return aggregator(name, combiner); + } + + private Aggregator aggregator( + final String name, + final CombineFn combiner) { + Aggregator aggregator = new Aggregator() { + @Override + public void addValue(AinT value) { + AccT accum = (AccT) accumulators.get(name); + AccT newAccum = combiner.addInput(accum, value); + accumulators.put(name, newAccum); + } + + @Override + public String getName() { + return name; + } + + @Override + public CombineFn getCombineFn() { + return combiner; + } + }; + accumulators.put(name, combiner.createAccumulator()); + return aggregator; + } + + @Override + public void sideOutputWithTimestamp(TupleTag tag, T output, Instant timestamp) { + + } + + @Override + public void sideOutput(TupleTag tag, T output) { + sideOutputWithTimestamp(tag, output, BoundedWindow.TIMESTAMP_MIN_VALUE); + } + + public void noteOutput(TupleTag tag, WindowedValue output) { + getOutputList(tag).add(output); + } + + private List> getOutputList(TupleTag tag) { + @SuppressWarnings({"unchecked", "rawtypes"}) + List> outputList = (List) outputs.get(tag); + if (outputList == null) { + outputList = new ArrayList<>(); + outputs.put(tag, (List) outputList); + } + return outputList; + } + } + + private TestProcessContext createProcessContext( + DoFn fn, + InputT elem) { + return new TestProcessContext<>(fn, + createContext(fn), + WindowedValue.valueInGlobalWindow(elem), + mainOutputTag, + sideInputs); + } + + private static class TestProcessContext extends DoFn.ProcessContext { + private final TestContext context; + private final TupleTag mainOutputTag; + private final WindowedValue element; + private final Map, ?> sideInputs; + + private TestProcessContext( + DoFn fn, + TestContext context, + WindowedValue element, + TupleTag mainOutputTag, + Map, ?> sideInputs) { + fn.super(); + this.context = context; + this.element = element; + this.mainOutputTag = mainOutputTag; + this.sideInputs = sideInputs; + } + + @Override + public InT element() { + return element.getValue(); + } + + @Override + public T sideInput(PCollectionView view) { + @SuppressWarnings("unchecked") + T sideInput = (T) sideInputs.get(view); + return sideInput; + } + + @Override + public Instant timestamp() { + return element.getTimestamp(); + } + + @Override + public BoundedWindow window() { + return Iterables.getOnlyElement(element.getWindows()); + } + + @Override + public PaneInfo pane() { + return element.getPane(); + } + + @Override + public WindowingInternals windowingInternals() { + throw new UnsupportedOperationException( + "WindowingInternals is an internal implementation detail of the Beam SDK, " + + "and should not be used by user code"); + } + + @Override + public PipelineOptions getPipelineOptions() { + return context.getPipelineOptions(); + } + + @Override + public void output(OutT output) { + sideOutput(mainOutputTag, output); + } + + @Override + public void outputWithTimestamp(OutT output, Instant timestamp) { + sideOutputWithTimestamp(mainOutputTag, output, timestamp); + } + + @Override + public void sideOutput(TupleTag tag, T output) { + sideOutputWithTimestamp(tag, output, element.getTimestamp()); + } + + @Override + public void sideOutputWithTimestamp(TupleTag tag, T output, Instant timestamp) { + context.noteOutput(tag, + WindowedValue.of(output, timestamp, element.getWindows(), element.getPane())); + } + + @Override + protected Aggregator createAggregatorInternal( + String name, CombineFn combiner) { + throw new IllegalStateException("Aggregators should not be created within ProcessContext. " + + "Instead, create an aggregator at DoFn construction time with createAggregator, and " + + "ensure they are set up by the time startBundle is called " + + "with setupDelegateAggregators."); + } + } + ///////////////////////////////////////////////////////////////////////////// /** The possible states of processing a DoFn. */ @@ -438,6 +652,8 @@ enum State { private Map, Iterable>> sideInputs = new HashMap<>(); + private Map accumulators; + /** The output tags used by the DoFn under test. */ TupleTag mainOutputTag = new TupleTag<>(); List> sideOutputTags = new ArrayList<>(); @@ -446,13 +662,7 @@ enum State { DoFn fn; /** The ListOutputManager to examine the outputs. */ - DoFnRunnerBase.ListOutputManager outputManager; - - /** The DoFnRunner if processing is in progress. */ - DoFnRunner fnRunner; - - /** Counters for user-defined Aggregators if processing is in progress. */ - CounterSet counterSet; + Map, List>> outputs; /** The state of processing of the DoFn under test. */ State state; @@ -464,9 +674,8 @@ enum State { void resetState() { fn = null; - outputManager = null; - fnRunner = null; - counterSet = null; + outputs = null; + accumulators = null; state = State.UNSTARTED; } @@ -476,23 +685,12 @@ void initializeState() { SerializableUtils.deserializeFromByteArray( SerializableUtils.serializeToByteArray(origFn), origFn.toString()); - counterSet = new CounterSet(); PTuple runnerSideInputs = PTuple.empty(); for (Map.Entry, Iterable>> entry : sideInputs.entrySet()) { runnerSideInputs = runnerSideInputs.and(entry.getKey().getTagInternal(), entry.getValue()); } - outputManager = new DoFnRunnerBase.ListOutputManager(); - fnRunner = - DoFnRunners.createDefault( - options, - fn, - DirectSideInputReader.of(runnerSideInputs), - outputManager, - mainOutputTag, - sideOutputTags, - DirectModeExecutionContext.create().getOrCreateStepContext(STEP_NAME, TRANSFORM_NAME), - counterSet.getAddCounterMutator(), - WindowingStrategy.globalDefault()); + outputs = new HashMap<>(); + accumulators = new HashMap<>(); } } From 3abf6491e5e1464b0039c7d4272d1d58ac72e513 Mon Sep 17 00:00:00 2001 From: Thomas Groh Date: Tue, 7 Jun 2016 14:07:53 -0700 Subject: [PATCH 2/3] fixup! Remove DirectRunner implementations from DoFnTester --- .../java/org/apache/beam/sdk/transforms/DoFnTester.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java index 569b86b87157..f6a343991acf 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java @@ -383,14 +383,13 @@ public List takeSideOutputElements(TupleTag tag) { * Returns the value of the provided {@link Aggregator}. */ public AggregateT getAggregatorValue(Aggregator agg) { - return extractAggregatorValue(agg, agg.getCombineFn()); + return extractAggregatorValue(agg.getName(), agg.getCombineFn()); } private AggregateT extractAggregatorValue( - Aggregator agg, - CombineFn combiner) { + String name, CombineFn combiner) { @SuppressWarnings("unchecked") - AccumT accumulator = (AccumT) accumulators.get(agg.getName()); + AccumT accumulator = (AccumT) accumulators.get(name); return combiner.extractOutput(accumulator); } From 914361096c6d09dfd1e5aacf1ccb71e92d61c9f3 Mon Sep 17 00:00:00 2001 From: Thomas Groh Date: Tue, 7 Jun 2016 15:43:04 -0700 Subject: [PATCH 3/3] fixup! Remove DirectRunner implementations from DoFnTester Throw exceptions explicitly --- .../apache/beam/examples/WordCountTest.java | 2 +- .../cookbook/CombinePerKeyExamplesTest.java | 4 +-- .../examples/cookbook/FilterExamplesTest.java | 4 +-- .../examples/cookbook/JoinExamplesTest.java | 4 +-- .../cookbook/MaxPerKeyExamplesTest.java | 4 +-- .../examples/cookbook/TriggerExampleTest.java | 2 +- .../examples/complete/game/UserScoreTest.java | 2 +- .../dataflow/DataflowPipelineRunnerTest.java | 31 +++++----------- .../beam/sdk/transforms/DoFnTester.java | 36 +++++++------------ .../beam/sdk/transforms/DoFnTesterTest.java | 16 ++++----- .../sdk/transforms/join/CoGroupByKeyTest.java | 2 +- .../src/test/java/WordCountTest.java | 2 +- 12 files changed, 43 insertions(+), 66 deletions(-) diff --git a/examples/java/src/test/java/org/apache/beam/examples/WordCountTest.java b/examples/java/src/test/java/org/apache/beam/examples/WordCountTest.java index c114d6e8b633..1382612bcfb4 100644 --- a/examples/java/src/test/java/org/apache/beam/examples/WordCountTest.java +++ b/examples/java/src/test/java/org/apache/beam/examples/WordCountTest.java @@ -48,7 +48,7 @@ public class WordCountTest { /** Example test that tests a specific DoFn. */ @Test - public void testExtractWordsFn() { + public void testExtractWordsFn() throws Exception { DoFnTester extractWordsFn = DoFnTester.of(new ExtractWordsFn()); diff --git a/examples/java/src/test/java/org/apache/beam/examples/cookbook/CombinePerKeyExamplesTest.java b/examples/java/src/test/java/org/apache/beam/examples/cookbook/CombinePerKeyExamplesTest.java index bb475c421010..5d6456a9233f 100644 --- a/examples/java/src/test/java/org/apache/beam/examples/cookbook/CombinePerKeyExamplesTest.java +++ b/examples/java/src/test/java/org/apache/beam/examples/cookbook/CombinePerKeyExamplesTest.java @@ -72,7 +72,7 @@ public class CombinePerKeyExamplesTest { .set("all_plays", "king_lear,macbeth"); @Test - public void testExtractLargeWordsFn() { + public void testExtractLargeWordsFn() throws Exception { DoFnTester> extractLargeWordsFn = DoFnTester.of(new ExtractLargeWordsFn()); List> results = extractLargeWordsFn.processBatch(ROWS_ARRAY); @@ -82,7 +82,7 @@ public void testExtractLargeWordsFn() { } @Test - public void testFormatShakespeareOutputFn() { + public void testFormatShakespeareOutputFn() throws Exception { DoFnTester, TableRow> formatShakespeareOutputFn = DoFnTester.of(new FormatShakespeareOutputFn()); List results = formatShakespeareOutputFn.processBatch(COMBINED_TUPLES_ARRAY); diff --git a/examples/java/src/test/java/org/apache/beam/examples/cookbook/FilterExamplesTest.java b/examples/java/src/test/java/org/apache/beam/examples/cookbook/FilterExamplesTest.java index 26393df11d40..13beab067380 100644 --- a/examples/java/src/test/java/org/apache/beam/examples/cookbook/FilterExamplesTest.java +++ b/examples/java/src/test/java/org/apache/beam/examples/cookbook/FilterExamplesTest.java @@ -68,7 +68,7 @@ public class FilterExamplesTest { @Test - public void testProjectionFn() { + public void testProjectionFn() throws Exception { DoFnTester projectionFn = DoFnTester.of(new ProjectionFn()); List results = projectionFn.processBatch(ROWS_ARRAY); @@ -78,7 +78,7 @@ public void testProjectionFn() { } @Test - public void testFilterSingleMonthDataFn() { + public void testFilterSingleMonthDataFn() throws Exception { DoFnTester filterSingleMonthDataFn = DoFnTester.of(new FilterSingleMonthDataFn(7)); List results = filterSingleMonthDataFn.processBatch(PROJROWS_ARRAY); diff --git a/examples/java/src/test/java/org/apache/beam/examples/cookbook/JoinExamplesTest.java b/examples/java/src/test/java/org/apache/beam/examples/cookbook/JoinExamplesTest.java index c7827d1cf2a9..22fe6a1fc59f 100644 --- a/examples/java/src/test/java/org/apache/beam/examples/cookbook/JoinExamplesTest.java +++ b/examples/java/src/test/java/org/apache/beam/examples/cookbook/JoinExamplesTest.java @@ -84,7 +84,7 @@ public class JoinExamplesTest { }; @Test - public void testExtractEventDataFn() { + public void testExtractEventDataFn() throws Exception { DoFnTester> extractEventDataFn = DoFnTester.of(new ExtractEventDataFn()); List> results = extractEventDataFn.processBatch(EVENTS); @@ -93,7 +93,7 @@ public void testExtractEventDataFn() { } @Test - public void testExtractCountryInfoFn() { + public void testExtractCountryInfoFn() throws Exception { DoFnTester> extractCountryInfoFn = DoFnTester.of(new ExtractCountryInfoFn()); List> results = extractCountryInfoFn.processBatch(CCS); diff --git a/examples/java/src/test/java/org/apache/beam/examples/cookbook/MaxPerKeyExamplesTest.java b/examples/java/src/test/java/org/apache/beam/examples/cookbook/MaxPerKeyExamplesTest.java index 0a8462134a01..9e129a15748d 100644 --- a/examples/java/src/test/java/org/apache/beam/examples/cookbook/MaxPerKeyExamplesTest.java +++ b/examples/java/src/test/java/org/apache/beam/examples/cookbook/MaxPerKeyExamplesTest.java @@ -66,7 +66,7 @@ public class MaxPerKeyExamplesTest { @Test - public void testExtractTempFn() { + public void testExtractTempFn() throws Exception { DoFnTester> extractTempFn = DoFnTester.of(new ExtractTempFn()); List> results = extractTempFn.processBatch(TEST_ROWS); @@ -76,7 +76,7 @@ public void testExtractTempFn() { } @Test - public void testFormatMaxesFn() { + public void testFormatMaxesFn() throws Exception { DoFnTester, TableRow> formatMaxesFnFn = DoFnTester.of(new FormatMaxesFn()); List results = formatMaxesFnFn.processBatch(TEST_KVS); diff --git a/examples/java/src/test/java/org/apache/beam/examples/cookbook/TriggerExampleTest.java b/examples/java/src/test/java/org/apache/beam/examples/cookbook/TriggerExampleTest.java index 4f15aef21a78..fe75d147c738 100644 --- a/examples/java/src/test/java/org/apache/beam/examples/cookbook/TriggerExampleTest.java +++ b/examples/java/src/test/java/org/apache/beam/examples/cookbook/TriggerExampleTest.java @@ -87,7 +87,7 @@ public class TriggerExampleTest { .set("window", "[1970-01-01T00:00:00.000Z..1970-01-01T00:01:00.000Z)"); @Test - public void testExtractTotalFlow() { + public void testExtractTotalFlow() throws Exception { DoFnTester> extractFlowInfow = DoFnTester .of(new ExtractFlowInfo()); diff --git a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java index 40740a6bc2a3..842ea30d15b1 100644 --- a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java +++ b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java @@ -85,7 +85,7 @@ public class UserScoreTest implements Serializable { /** Test the ParseEventFn DoFn. */ @Test - public void testParseEventFn() { + public void testParseEventFn() throws Exception { DoFnTester parseEventFn = DoFnTester.of(new ParseEventFn()); diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineRunnerTest.java index 66c2feb53072..aa65dd1939d7 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineRunnerTest.java @@ -18,7 +18,6 @@ package org.apache.beam.runners.dataflow; import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow; - import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; @@ -76,7 +75,6 @@ import org.apache.beam.sdk.util.NoopPathValidator; import org.apache.beam.sdk.util.ReleaseInfo; import org.apache.beam.sdk.util.TestCredential; -import org.apache.beam.sdk.util.UserCodeException; import org.apache.beam.sdk.util.WindowedValue; import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder; import org.apache.beam.sdk.util.WindowingStrategy; @@ -976,18 +974,12 @@ public void testBatchViewAsSingletonToIsmRecordWithMultipleValuesThrowsException new BatchViewAsSingleton.IsmRecordForSingularValuePerWindowDoFn (GlobalWindow.Coder.INSTANCE)); - try { - doFnTester.processBatch( - ImmutableList.of(KV.>>>of( - 0, ImmutableList.of( - KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("a")), - KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("b")))))); - fail("Expected UserCodeException"); - } catch (UserCodeException e) { - assertTrue(e.getCause() instanceof IllegalStateException); - IllegalStateException rootCause = (IllegalStateException) e.getCause(); - assertThat(rootCause.getMessage(), containsString("found for singleton within window")); - } + thrown.expect(IllegalStateException.class); + thrown.expectMessage("found for singleton within window"); + doFnTester.processBatch(ImmutableList.of( + KV.>>>of(0, + ImmutableList.of(KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("a")), + KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("b")))))); } @Test @@ -1192,14 +1184,9 @@ public void testToIsmRecordForMapLikeDoFnWithoutUniqueKeysThrowsException() thro KV.of(KV.of(1L, windowA), WindowedValue.of(111L, new Instant(2), windowA, PaneInfo.NO_FIRING))))); - try { - doFnTester.processBatch(inputElements); - fail("Expected UserCodeException"); - } catch (UserCodeException e) { - assertTrue(e.getCause() instanceof IllegalStateException); - IllegalStateException rootCause = (IllegalStateException) e.getCause(); - assertThat(rootCause.getMessage(), containsString("Unique keys are expected but found key")); - } + thrown.expect(IllegalStateException.class); + thrown.expectMessage("Unique keys are expected but found key"); + doFnTester.processBatch(inputElements); } @Test diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java index f6a343991acf..332ea1380cfd 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java @@ -25,7 +25,6 @@ import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.util.PTuple; import org.apache.beam.sdk.util.SerializableUtils; -import org.apache.beam.sdk.util.UserCodeException; import org.apache.beam.sdk.util.WindowedValue; import org.apache.beam.sdk.util.WindowingInternals; import org.apache.beam.sdk.values.PCollectionView; @@ -164,7 +163,7 @@ public void setSideOutputTags(TupleTagList sideOutputTags) { * calls {@link #finishBundle}, then returns the result of * {@link #takeOutputElements}. */ - public List processBatch(Iterable inputElements) { + public List processBatch(Iterable inputElements) throws Exception { startBundle(); for (InputT inputElement : inputElements) { processElement(inputElement); @@ -185,7 +184,7 @@ public List processBatch(Iterable inputElements) { * */ @SafeVarargs - public final List processBatch(InputT... inputElements) { + public final List processBatch(InputT... inputElements) throws Exception { return processBatch(Arrays.asList(inputElements)); } @@ -194,16 +193,12 @@ public final List processBatch(InputT... inputElements) { * *

If needed, first creates a fresh instance of the DoFn under test. */ - public void startBundle() { + public void startBundle() throws Exception { resetState(); initializeState(); - try { - TestContext context = createContext(fn); - context.setupDelegateAggregators(); - fn.startBundle(context); - } catch (Exception e) { - throw UserCodeException.wrap(e); - } + TestContext context = createContext(fn); + context.setupDelegateAggregators(); + fn.startBundle(context); state = State.STARTED; } @@ -218,18 +213,14 @@ public void startBundle() { * @throws IllegalStateException if the {@code DoFn} under test has already * been finished */ - public void processElement(InputT element) { + public void processElement(InputT element) throws Exception { if (state == State.FINISHED) { throw new IllegalStateException("finishBundle() has already been called"); } if (state == State.UNSTARTED) { startBundle(); } - try { - fn.processElement(createProcessContext(fn, element)); - } catch (Exception e) { - throw UserCodeException.wrap(e); - } + fn.processElement(createProcessContext(fn, element)); } /** @@ -241,18 +232,14 @@ public void processElement(InputT element) { * @throws IllegalStateException if the {@code DoFn} under test has already * been finished */ - public void finishBundle() { + public void finishBundle() throws Exception { if (state == State.FINISHED) { throw new IllegalStateException("finishBundle() has already been called"); } if (state == State.UNSTARTED) { startBundle(); } - try { - fn.finishBundle(createContext(fn)); - } catch (Exception e) { - throw UserCodeException.wrap(e); - } + fn.finishBundle(createContext(fn)); state = State.FINISHED; } @@ -390,6 +377,9 @@ private AggregateT extractAggregatorValue( String name, CombineFn combiner) { @SuppressWarnings("unchecked") AccumT accumulator = (AccumT) accumulators.get(name); + if (accumulator == null) { + accumulator = combiner.createAccumulator(); + } return combiner.extractOutput(accumulator); } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java index 4df5be7366c1..ec22251b671d 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java @@ -17,8 +17,8 @@ */ package org.apache.beam.sdk.transforms; -import static org.hamcrest.CoreMatchers.hasItems; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.hasItems; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThat; import static org.junit.Assert.assertTrue; @@ -39,7 +39,7 @@ public class DoFnTesterTest { @Test - public void processElement() { + public void processElement() throws Exception { CounterDoFn counterDoFn = new CounterDoFn(); DoFnTester tester = DoFnTester.of(counterDoFn); @@ -61,7 +61,7 @@ public void processElement() { } @Test - public void processElementsWithPeeks() { + public void processElementsWithPeeks() throws Exception { CounterDoFn counterDoFn = new CounterDoFn(); DoFnTester tester = DoFnTester.of(counterDoFn); @@ -119,7 +119,7 @@ public void processElementsWithPeeks() { } @Test - public void processBatch() { + public void processBatch() throws Exception { CounterDoFn counterDoFn = new CounterDoFn(); DoFnTester tester = DoFnTester.of(counterDoFn); @@ -138,7 +138,7 @@ public void processBatch() { } @Test - public void processElementWithTimestamp() { + public void processElementWithTimestamp() throws Exception { CounterDoFn counterDoFn = new CounterDoFn(); DoFnTester tester = DoFnTester.of(counterDoFn); @@ -175,7 +175,7 @@ public void processElementWithTimestamp() { } @Test - public void getAggregatorValuesShouldGetValueOfCounter() { + public void getAggregatorValuesShouldGetValueOfCounter() throws Exception { CounterDoFn counterDoFn = new CounterDoFn(); DoFnTester tester = DoFnTester.of(counterDoFn); tester.processBatch(1L, 2L, 4L, 8L); @@ -186,7 +186,7 @@ public void getAggregatorValuesShouldGetValueOfCounter() { } @Test - public void getAggregatorValuesWithEmptyCounterShouldSucceed() { + public void getAggregatorValuesWithEmptyCounterShouldSucceed() throws Exception { CounterDoFn counterDoFn = new CounterDoFn(); DoFnTester tester = DoFnTester.of(counterDoFn); tester.processBatch(); @@ -196,7 +196,7 @@ public void getAggregatorValuesWithEmptyCounterShouldSucceed() { } @Test - public void getAggregatorValuesInStartFinishBundleShouldGetValues() { + public void getAggregatorValuesInStartFinishBundleShouldGetValues() throws Exception { CounterDoFn fn = new CounterDoFn(1L, 2L); DoFnTester tester = DoFnTester.of(fn); tester.processBatch(0L, 0L); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/CoGroupByKeyTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/CoGroupByKeyTest.java index f8e7f08e4c7b..d99c536ca466 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/CoGroupByKeyTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/CoGroupByKeyTest.java @@ -406,7 +406,7 @@ public void processElement(ProcessContext c) { */ @SuppressWarnings("unchecked") @Test - public void testConsumingDoFn() { + public void testConsumingDoFn() throws Exception { TupleTag purchasesTag = new TupleTag<>(); TupleTag addressesTag = new TupleTag<>(); TupleTag namesTag = new TupleTag<>(); diff --git a/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/src/test/java/WordCountTest.java b/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/src/test/java/WordCountTest.java index b921875e0cbf..4b2619874029 100644 --- a/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/src/test/java/WordCountTest.java +++ b/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/src/test/java/WordCountTest.java @@ -48,7 +48,7 @@ public class WordCountTest { /** Example test that tests a specific DoFn. */ @Test - public void testExtractWordsFn() { + public void testExtractWordsFn() throws Exception { DoFnTester extractWordsFn = DoFnTester.of(new ExtractWordsFn());