org.apache.beam.sdk.transforms

Class DoFnTester<InputT,OutputT>

java.lang.Object

org.apache.beam.sdk.transforms.DoFnTester<InputT,OutputT>

Type Parameters:

InputT - the type of the DoFn's (main) input elements

OutputT - the type of the DoFn's (main) output elements

All Implemented Interfaces:

java.lang.AutoCloseable

public class DoFnTester<InputT,OutputT>
extends java.lang.Object
implements java.lang.AutoCloseable

A harness for unit-testing a DoFn.

For example:

 
 DoFn<InputT, OutputT> fn = ...;

 DoFnTester<InputT, OutputT> fnTester = DoFnTester.of(fn);

 // Set arguments shared across all bundles:
 fnTester.setSideInputs(...);      // If fn takes side inputs.
 fnTester.setOutputTags(...);  // If fn writes to more than one output.

 // Process a bundle containing a single input element:
 Input testInput = ...;
 List<OutputT> testOutputs = fnTester.processBundle(testInput);
 Assert.assertThat(testOutputs, Matchers.hasItems(...));

 // Process a bigger bundle:
 Assert.assertThat(fnTester.processBundle(i1, i2, ...), Matchers.hasItems(...));
  

Nested Class Summary

Nested Classes

Modifier and Type	Class and Description
static class	DoFnTester.CloningBehavior When a DoFnTester should clone the DoFn under test and how it should manage the lifecycle of the DoFn.

Method Summary

Modifier and Type	Method and Description
void	clearOutputElements() Clears the record of the elements output so far to the main output.
<T> void	clearOutputElements(TupleTag<T> tag) Clears the record of the elements output so far to the output with the given tag.
void	close()
DoFn.ProcessContext	createProcessContext(ValueInSingleWindow<InputT> element)
void	finishBundle() Calls the DoFn.FinishBundle method of the DoFn under test.
DoFnTester.CloningBehavior	getCloningBehavior() Indicates whether this DoFnTester will clone the DoFn under test.
TupleTag<OutputT>	getMainOutputTag()
<T> java.util.List<ValueInSingleWindow<T>>	getMutableOutput(TupleTag<T> tag)
PipelineOptions	getPipelineOptions()
static <InputT,OutputT> DoFnTester<InputT,OutputT>	of(DoFn<InputT,OutputT> fn) Returns a DoFnTester supporting unit-testing of the given DoFn.
java.util.List<OutputT>	peekOutputElements() Returns the elements output so far to the main output.
<T> java.util.List<T>	peekOutputElements(TupleTag<T> tag) Returns the elements output so far to the output with the given tag.
java.util.List<TimestampedValue<OutputT>>	peekOutputElementsInWindow(BoundedWindow window) Returns the elements output so far to the main output in the provided window with associated timestamps.
java.util.List<TimestampedValue<OutputT>>	peekOutputElementsInWindow(TupleTag<OutputT> tag, BoundedWindow window) Returns the elements output so far to the specified output in the provided window with associated timestamps.
java.util.List<TimestampedValue<OutputT>>	peekOutputElementsWithTimestamp() Returns the elements output so far to the main output with associated timestamps.
java.util.List<OutputT>	processBundle(InputT... inputElements) A convenience method for testing DoFns with bundles of elements.
java.util.List<OutputT>	processBundle(java.lang.Iterable<? extends InputT> inputElements) A convenience operation that first calls startBundle(), then calls processElement(InputT) on each of the input elements, then calls finishBundle(), then returns the result of takeOutputElements().
void	processElement(InputT element) Calls the DoFn.ProcessElement method on the DoFn under test, in a context where DoFn.ProcessContext.element() returns the given element and the element is in the global window.
void	processTimestampedElement(TimestampedValue<InputT> element) Calls DoFn.ProcessElement on the DoFn under test, in a context where DoFn.ProcessContext.element() returns the given element and timestamp and the element is in the global window.
void	processWindowedElement(InputT element, Instant timestamp, BoundedWindow window) Calls DoFn.ProcessElement on the DoFn under test, in a context where DoFn.ProcessContext.element() returns the given element and timestamp and the element is in the given window.
void	setCloningBehavior(DoFnTester.CloningBehavior newValue) Instruct this DoFnTester whether or not to clone the DoFn under test.
<T> void	setSideInput(PCollectionView<T> sideInput, BoundedWindow window, T value) Registers the values of a side input PCollectionView to pass to the DoFn under test.
void	setSideInputs(java.util.Map<PCollectionView<?>,java.util.Map<BoundedWindow,?>> sideInputs) Registers the tuple of values of the side input PCollectionViews to pass to the DoFn under test.
void	startBundle() Calls the DoFn.StartBundle method on the DoFn under test.
java.util.List<OutputT>	takeOutputElements() Returns the elements output so far to the main output.
<T> java.util.List<T>	takeOutputElements(TupleTag<T> tag) Returns the elements output so far to the output with the given tag.
java.util.List<TimestampedValue<OutputT>>	takeOutputElementsWithTimestamp() Returns the elements output so far to the main output with associated timestamps.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Method Detail

of

public static <InputT,OutputT> DoFnTester<InputT,OutputT> of(DoFn<InputT,OutputT> fn)

Returns a DoFnTester supporting unit-testing of the given DoFn. By default, uses DoFnTester.CloningBehavior.CLONE_ONCE.

The only supported extra parameter of the DoFn.ProcessElement method is BoundedWindow.

setSideInputs

public void setSideInputs(java.util.Map<PCollectionView<?>,java.util.Map<BoundedWindow,?>> sideInputs)

Registers the tuple of values of the side input PCollectionViews to pass to the DoFn under test.

Resets the state of this DoFnTester.

If this isn't called, DoFnTester assumes the DoFn takes no side inputs.

setSideInput

public <T> void setSideInput(PCollectionView<T> sideInput,
                             BoundedWindow window,
                             T value)

Registers the values of a side input PCollectionView to pass to the DoFn under test.

The provided value is the final value of the side input in the specified window, not the value of the input PCollection in that window.

If this isn't called, DoFnTester will return the default value for any side input that is used.

getPipelineOptions

public PipelineOptions getPipelineOptions()

setCloningBehavior

public void setCloningBehavior(DoFnTester.CloningBehavior newValue)

Instruct this DoFnTester whether or not to clone the DoFn under test.

getCloningBehavior

public DoFnTester.CloningBehavior getCloningBehavior()

Indicates whether this DoFnTester will clone the DoFn under test.

processBundle

public java.util.List<OutputT> processBundle(java.lang.Iterable<? extends InputT> inputElements)
                                      throws java.lang.Exception

A convenience operation that first calls startBundle(), then calls processElement(InputT) on each of the input elements, then calls finishBundle(), then returns the result of takeOutputElements().

Throws:

java.lang.Exception

processBundle

@SafeVarargs
public final java.util.List<OutputT> processBundle(InputT... inputElements)
                                                         throws java.lang.Exception

A convenience method for testing DoFns with bundles of elements. Logic proceeds as follows:

1. Calls startBundle().
2. Calls processElement(InputT) on each of the arguments.
3. Calls finishBundle().
4. Returns the result of takeOutputElements().

Throws:

java.lang.Exception

startBundle

public void startBundle()
                 throws java.lang.Exception

Calls the DoFn.StartBundle method on the DoFn under test.

If needed, first creates a fresh instance of the DoFn under test and calls DoFn.Setup.

Throws:

java.lang.Exception

processElement

public void processElement(InputT element)
                    throws java.lang.Exception

Calls the DoFn.ProcessElement method on the DoFn under test, in a context where DoFn.ProcessContext.element() returns the given element and the element is in the global window.

Will call startBundle() automatically, if it hasn't already been called.

Throws:

java.lang.IllegalStateException - if the DoFn under test has already been finished

java.lang.Exception

processTimestampedElement

public void processTimestampedElement(TimestampedValue<InputT> element)
                               throws java.lang.Exception

Calls DoFn.ProcessElement on the DoFn under test, in a context where DoFn.ProcessContext.element() returns the given element and timestamp and the element is in the global window.

Will call startBundle() automatically, if it hasn't already been called.

Throws:

java.lang.Exception

processWindowedElement

public void processWindowedElement(InputT element,
                                   Instant timestamp,
                                   BoundedWindow window)
                            throws java.lang.Exception

Calls DoFn.ProcessElement on the DoFn under test, in a context where DoFn.ProcessContext.element() returns the given element and timestamp and the element is in the given window.

Will call startBundle() automatically, if it hasn't already been called.

Throws:

java.lang.Exception

finishBundle

public void finishBundle()
                  throws java.lang.Exception

Calls the DoFn.FinishBundle method of the DoFn under test.

If setCloningBehavior(org.apache.beam.sdk.transforms.DoFnTester.CloningBehavior) was called with DoFnTester.CloningBehavior.CLONE_PER_BUNDLE, then also calls DoFn.Teardown on the DoFn, and it will be cloned and DoFn.Setup again when processing the next bundle.

Throws:

java.lang.IllegalStateException - if DoFn.FinishBundle has already been called for this bundle.

java.lang.Exception

peekOutputElements

public java.util.List<OutputT> peekOutputElements()

Returns the elements output so far to the main output. Does not clear them, so subsequent calls will continue to include these elements.

See Also:

takeOutputElements(), clearOutputElements()

peekOutputElementsWithTimestamp

@Experimental
public java.util.List<TimestampedValue<OutputT>> peekOutputElementsWithTimestamp()

Returns the elements output so far to the main output with associated timestamps. Does not clear them, so subsequent calls will continue to include these. elements.

See Also:

takeOutputElementsWithTimestamp(), clearOutputElements()

peekOutputElementsInWindow

public java.util.List<TimestampedValue<OutputT>> peekOutputElementsInWindow(BoundedWindow window)

Returns the elements output so far to the main output in the provided window with associated timestamps.

peekOutputElementsInWindow

public java.util.List<TimestampedValue<OutputT>> peekOutputElementsInWindow(TupleTag<OutputT> tag,
                                                                            BoundedWindow window)

Returns the elements output so far to the specified output in the provided window with associated timestamps.

clearOutputElements

public void clearOutputElements()

Clears the record of the elements output so far to the main output.

See Also:

peekOutputElements()

takeOutputElements

public java.util.List<OutputT> takeOutputElements()

Returns the elements output so far to the main output. Clears the list so these elements don't appear in future calls.

See Also:

peekOutputElements()

takeOutputElementsWithTimestamp

@Experimental
public java.util.List<TimestampedValue<OutputT>> takeOutputElementsWithTimestamp()

Returns the elements output so far to the main output with associated timestamps. Clears the list so these elements don't appear in future calls.

See Also:

peekOutputElementsWithTimestamp(), takeOutputElements(), clearOutputElements()

peekOutputElements

public <T> java.util.List<T> peekOutputElements(TupleTag<T> tag)

Returns the elements output so far to the output with the given tag. Does not clear them, so subsequent calls will continue to include these elements.

See Also:

takeOutputElements(), clearOutputElements()

clearOutputElements

public <T> void clearOutputElements(TupleTag<T> tag)

Clears the record of the elements output so far to the output with the given tag.

See Also:

peekOutputElements()

takeOutputElements

public <T> java.util.List<T> takeOutputElements(TupleTag<T> tag)

Returns the elements output so far to the output with the given tag. Clears the list so these elements don't appear in future calls.

See Also:

peekOutputElements()

getMutableOutput

public <T> java.util.List<ValueInSingleWindow<T>> getMutableOutput(TupleTag<T> tag)

getMainOutputTag

public TupleTag<OutputT> getMainOutputTag()

createProcessContext

public DoFn.ProcessContext createProcessContext(ValueInSingleWindow<InputT> element)

close

public void close()
           throws java.lang.Exception

Specified by:

close in interface java.lang.AutoCloseable

Throws:

java.lang.Exception