Switch DoFnInvokersTest to use a mock ExtraContextFactory
Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/c052d2a1 Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/c052d2a1 Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/c052d2a1 Branch: refs/heads/master Commit: c052d2a1a1fd36f2c21e63427a9e4a50addf85f7 Parents: 40ff9d4 Author: Kenneth Knowles <k...@google.com> Authored: Thu Nov 3 19:26:36 2016 -0700 Committer: Kenneth Knowles <k...@google.com> Committed: Mon Nov 7 15:25:03 2016 -0800 ---------------------------------------------------------------------- .../transforms/reflect/DoFnInvokersTest.java | 35 ++++---------------- 1 file changed, 6 insertions(+), 29 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/c052d2a1/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java ---------------------------------------------------------------------- diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java index 0bfe2be..dbb7955 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java @@ -39,6 +39,7 @@ import org.apache.beam.sdk.coders.CoderRegistry; import org.apache.beam.sdk.coders.CustomCoder; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.DoFn.GetInitialRestriction; +import org.apache.beam.sdk.transforms.DoFn.ExtraContextFactory; import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation; import org.apache.beam.sdk.transforms.OldDoFn; import org.apache.beam.sdk.transforms.reflect.testhelper.DoFnInvokersTestHelper; @@ -67,41 +68,17 @@ public class DoFnInvokersTest { @Mock private DoFn.InputProvider<String> mockInputProvider; @Mock private DoFn.OutputReceiver<String> mockOutputReceiver; @Mock private WindowingInternals<String, String> mockWindowingInternals; + @Mock private ExtraContextFactory<String, String> extraContextFactory; @Mock private OldDoFn<String, String> mockOldDoFn; - private DoFn.ExtraContextFactory<String, String> extraContextFactory; - @Before public void setUp() { MockitoAnnotations.initMocks(this); - this.extraContextFactory = - new DoFn.ExtraContextFactory<String, String>() { - @Override - public BoundedWindow window() { - return mockWindow; - } - - @Override - public DoFn.InputProvider<String> inputProvider() { - return mockInputProvider; - } - - @Override - public DoFn.OutputReceiver<String> outputReceiver() { - return mockOutputReceiver; - } - - @Override - public WindowingInternals<String, String> windowingInternals() { - return mockWindowingInternals; - } - - @Override - public <RestrictionT> RestrictionTracker<RestrictionT> restrictionTracker() { - return null; - } - }; + when(extraContextFactory.window()).thenReturn(mockWindow); + when(extraContextFactory.inputProvider()).thenReturn(mockInputProvider); + when(extraContextFactory.outputReceiver()).thenReturn(mockOutputReceiver); + when(extraContextFactory.windowingInternals()).thenReturn(mockWindowingInternals); } private ProcessContinuation invokeProcessElement(DoFn<String, String> fn) {