boyuanzz commented on a change in pull request #12016:
URL: https://github.com/apache/beam/pull/12016#discussion_r454081616
##########
File path:
sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
##########
@@ -742,10 +785,202 @@ public void outputWithTimestamp(String output, Instant
instant) {
new FakeArgumentProvider<String, String>() {
@Override
public Object restriction() {
- return new RestrictionWithDefaultTracker();
+ return new RestrictionWithBoundedDefaultTracker();
}
}),
- instanceOf(DefaultTracker.class));
+ instanceOf(BoundedDefaultTracker.class));
+ assertThat(
+ invoker.invokeNewWatermarkEstimator(
+ new FakeArgumentProvider<String, String>() {
+ @Override
+ public Object watermarkEstimatorState() {
+ return new
WatermarkEstimatorStateWithDefaultWatermarkEstimator();
+ }
+ }),
+ instanceOf(DefaultWatermarkEstimator.class));
+ }
+
+ @Test
+ public void testTruncateFnWithHasDefaultMethodsWhenBounded() throws
Exception {
+ class BoundedMockFn extends DoFn<String, String> {
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void>
tracker,
+
WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator>
+ watermarkEstimator) {}
+
+ @GetInitialRestriction
+ public RestrictionWithBoundedDefaultTracker
getInitialRestriction(@Element String element) {
+ return null;
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public WatermarkEstimatorStateWithDefaultWatermarkEstimator
+ getInitialWatermarkEstimatorState() {
+ return null;
+ }
+ }
+
+ BoundedMockFn fn = mock(BoundedMockFn.class);
+ DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
+
+ CoderRegistry coderRegistry = CoderRegistry.createDefault();
+ coderRegistry.registerCoderProvider(
+ CoderProviders.fromStaticMethods(
+ RestrictionWithBoundedDefaultTracker.class,
CoderForDefaultTracker.class));
+ coderRegistry.registerCoderForClass(
+ WatermarkEstimatorStateWithDefaultWatermarkEstimator.class,
+ new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator());
+ assertThat(
+
invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry),
+ instanceOf(CoderForDefaultTracker.class));
+ assertThat(
+ invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry),
+
instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class));
+ RestrictionTracker tracker =
+ invoker.invokeNewTracker(
+ new FakeArgumentProvider<String, String>() {
+ @Override
+ public Object restriction() {
+ return new RestrictionWithBoundedDefaultTracker();
+ }
+ });
+ assertThat(tracker, instanceOf(BoundedDefaultTracker.class));
+ invoker.invokeTruncateRestriction(
+ new FakeArgumentProvider<String, String>() {
+ @Override
+ public RestrictionTracker restrictionTracker() {
+ return tracker;
+ }
+
+ @Override
+ public String element(DoFn<String, String> doFn) {
+ return "blah";
+ }
+
+ @Override
+ public Object restriction() {
+ return "foo";
+ }
+
+ @Override
+ public OutputReceiver<String> outputReceiver(DoFn<String, String>
doFn) {
+ return new DoFn.OutputReceiver<String>() {
+ private boolean invoked;
+
+ @Override
+ public void output(String output) {
+ assertFalse(invoked);
+ invoked = true;
+ assertEquals("foo", output);
+ }
+
+ @Override
+ public void outputWithTimestamp(String output, Instant instant) {
+ assertFalse(invoked);
+ invoked = true;
+ assertEquals("foo", output);
+ }
+ };
+ }
+ });
+ assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider));
+ assertThat(
+ invoker.invokeNewWatermarkEstimator(
+ new FakeArgumentProvider<String, String>() {
+ @Override
+ public Object watermarkEstimatorState() {
+ return new
WatermarkEstimatorStateWithDefaultWatermarkEstimator();
+ }
+ }),
+ instanceOf(DefaultWatermarkEstimator.class));
+ }
+
+ @Test
+ public void testTruncateFnWithHasDefaultMethodsWhenUnbounded() throws
Exception {
+ class UnboundedMockFn extends DoFn<String, String> {
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ RestrictionTracker<RestrictionWithUnboundedDefaultTracker, Void>
tracker,
+
WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator>
+ watermarkEstimator) {}
+
+ @GetInitialRestriction
+ public RestrictionWithUnboundedDefaultTracker
getInitialRestriction(@Element String element) {
+ return null;
+ }
+
+ @GetInitialWatermarkEstimatorState
+ public WatermarkEstimatorStateWithDefaultWatermarkEstimator
+ getInitialWatermarkEstimatorState() {
+ return null;
+ }
+ }
+
+ UnboundedMockFn fn = mock(UnboundedMockFn.class);
+ DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
+
+ CoderRegistry coderRegistry = CoderRegistry.createDefault();
+ coderRegistry.registerCoderProvider(
+ CoderProviders.fromStaticMethods(
+ RestrictionWithUnboundedDefaultTracker.class,
CoderForDefaultTracker.class));
+ coderRegistry.registerCoderForClass(
+ WatermarkEstimatorStateWithDefaultWatermarkEstimator.class,
+ new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator());
+ assertThat(
+
invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry),
+ instanceOf(CoderForDefaultTracker.class));
+ assertThat(
+ invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry),
+
instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class));
+ RestrictionTracker tracker =
+ invoker.invokeNewTracker(
+ new FakeArgumentProvider<String, String>() {
+ @Override
+ public Object restriction() {
+ return new RestrictionWithUnboundedDefaultTracker();
+ }
+ });
+ assertThat(tracker, instanceOf(UnboundedDefaultTracker.class));
+ invoker.invokeTruncateRestriction(
+ new FakeArgumentProvider<String, String>() {
+ @Override
+ public RestrictionTracker restrictionTracker() {
+ return tracker;
+ }
+
+ @Override
+ public String element(DoFn<String, String> doFn) {
+ return "blah";
+ }
+
+ @Override
+ public Object restriction() {
+ return "foo";
+ }
+
+ @Override
+ public OutputReceiver<String> outputReceiver(DoFn<String, String>
doFn) {
+ return new DoFn.OutputReceiver<String>() {
+ private final boolean shouldInvoked = false;
+
+ // This should not be invoked.
+ @Override
+ public void output(String output) {
+ assertTrue(shouldInvoked);
+ }
+
+ // This should not be invoked.
+ @Override
+ public void outputWithTimestamp(String output, Instant instant) {
+ assertTrue(shouldInvoked);
+ }
Review comment:
It's no longer applied.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
