boyuanzz commented on a change in pull request #12016:
URL: https://github.com/apache/beam/pull/12016#discussion_r454081206
##########
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);
Review comment:
Given that we decided to make `@TruncateRestriction` returns truncated
restrictions, changes here are no longer applied. Do you want me to create a
JIRA for other tests?
----------------------------------------------------------------
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]
