PTransformMatchers { public static PTransformMatcher parDoWithFnType(final Class<? extends DoFn> fnType) { return new PTransformMatcher() { @Override public boolean matches(AppliedPTransform<?, ?, ?> application) { DoFn<?, ?> fn; if (application.getTransform() instanceof ParDo.SingleOutput) { fn = ((ParDo.SingleOutput) application.getTransform()).getFn(); } else if (application.getTransform() instanceof ParDo.MultiOutput) { fn = ((ParDo.MultiOutput) application.getTransform()).getFn(); } else { return false; } return fnType.equals(fn.getClass()); } @Override public String toString() { return MoreObjects.toStringHelper("ParDoWithFnTypeMatcher") .add("fnType", fnType) .toString(); } }; } private  PTransformMatchers(); static PTransformMatcher urnEqualTo(String urn); static PTransformMatcher classEqualTo(Class<? extends PTransform> clazz); static PTransformMatcher requiresStableInputParDoSingle(); static PTransformMatcher requiresStableInputParDoMulti(); static PTransformMatcher splittableParDoSingle(); static PTransformMatcher stateOrTimerParDoSingle(); static PTransformMatcher splittableParDoMulti(); static PTransformMatcher splittableParDo(); static PTransformMatcher splittableProcessKeyedBounded(); static PTransformMatcher splittableProcessKeyedUnbounded(); static PTransformMatcher stateOrTimerParDo(); static PTransformMatcher stateOrTimerParDoMulti(); static PTransformMatcher parDoWithFnType(final Class<? extends DoFn> fnType); static PTransformMatcher createViewWithViewFn(final Class<? extends ViewFn> viewFnType); static PTransformMatcher emptyFlatten(); static PTransformMatcher flattenWithDuplicateInputs(); static PTransformMatcher groupWithShardableStates(); static PTransformMatcher writeWithRunnerDeterminedSharding();  }
@Test public void parDoWithFnTypeWithMatchingType() { DoFn<Object, Object> fn = new DoFn<Object, Object>() { @ProcessElement public void process(ProcessContext ctxt) {} }; AppliedPTransform<?, ?, ?> parDoSingle = getAppliedTransform(ParDo.of(fn)); AppliedPTransform<?, ?, ?> parDoMulti = getAppliedTransform(ParDo.of(fn).withOutputTags(new TupleTag<>(), TupleTagList.empty())); PTransformMatcher matcher = PTransformMatchers.parDoWithFnType(fn.getClass()); assertThat(matcher.matches(parDoSingle), is(true)); assertThat(matcher.matches(parDoMulti), is(true)); }
@Test public void parDoWithFnTypeWithNoMatch() { DoFn<Object, Object> fn = new DoFn<Object, Object>() { @ProcessElement public void process(ProcessContext ctxt) {} }; AppliedPTransform<?, ?, ?> parDoSingle = getAppliedTransform(ParDo.of(fn)); AppliedPTransform<?, ?, ?> parDoMulti = getAppliedTransform(ParDo.of(fn).withOutputTags(new TupleTag<>(), TupleTagList.empty())); PTransformMatcher matcher = PTransformMatchers.parDoWithFnType(doFnWithState.getClass()); assertThat(matcher.matches(parDoSingle), is(false)); assertThat(matcher.matches(parDoMulti), is(false)); }
@Test public void parDoWithFnTypeNotParDo() { AppliedPTransform<?, ?, ?> notParDo = getAppliedTransform(Create.empty(VoidCoder.of())); PTransformMatcher matcher = PTransformMatchers.parDoWithFnType(doFnWithState.getClass()); assertThat(matcher.matches(notParDo), is(false)); }