tests, remove default implementations.

Author: stankiewicz

runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java

@@ -555,6 +555,11 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
       return timestamp();
     }
 
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      return elem.causedByDrain();
+    }
+
     @Override
     public String timerId(DoFn<InputT, OutputT> doFn) {
       throw new UnsupportedOperationException(
@@ -831,6 +836,11 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
       return timestamp();
     }
 
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      return causedByDrain;
+    }
+
     @Override
     public String timerId(DoFn<InputT, OutputT> doFn) {
       return timerId;
@@ -1119,6 +1129,11 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
       return timestamp;
     }
 
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      throw new UnsupportedOperationException("CausedByDrain parameters are not supported.");
+    }
+
     @Override
     public String timerId(DoFn<InputT, OutputT> doFn) {
       throw new UnsupportedOperationException("Timer parameters are not supported.");

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java

@@ -219,9 +219,7 @@ interface ArgumentProvider<InputT, OutputT> {
     Instant timestamp(DoFn<InputT, OutputT> doFn);
 
     /** Provide a reference to the caused by drain. */
-    default CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
-      return CausedByDrain.NORMAL;
-    }
+    CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn);
 
     /** Provide a reference to the time domain for a timer firing. */
     TimeDomain timeDomain(DoFn<InputT, OutputT> doFn);

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java

@@ -1362,7 +1362,10 @@ private static Parameter analyzeExtraParameter(
       return Parameter.keyT(paramT);
     } else if (rawType.equals(TimeDomain.class)) {
       return Parameter.timeDomainParameter();
-    } else if (rawType.equals(CausedByDrain.class)) {
+    } else if (CausedByDrain.class.isAssignableFrom(rawType)) {
+      methodErrors.checkArgument(
+          rawType.equals(CausedByDrain.class),
+          "CausedByDrain argument must have type org.apache.beam.sdk.values.CausedByDrain.");
       return Parameter.causedByDrainParameter();
     } else if (hasAnnotation(DoFn.SideInput.class, param.getAnnotations())) {
       String sideInputId = getSideInputId(param.getAnnotations());

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java

@@ -543,6 +543,11 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
         return outerContext.timestamp();
       }
 
+      @Override
+      public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+        return outerContext.causedByDrain();
+      }
+
       @Override
       public String timerId(DoFn<InputT, OutputT> doFn) {
         throw new UnsupportedOperationException();

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java

@@ -56,6 +56,7 @@
 import org.apache.beam.sdk.transforms.Sum;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.BundleFinalizerParameter;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.CausedByDrainParameter;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.ElementParameter;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.FinishBundleContextParameter;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.OutputReceiverParameter;
@@ -78,6 +79,7 @@
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
+import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TypeDescriptor;
@@ -130,10 +132,11 @@ public void process(
                   PipelineOptions options,
                   @SideInput("tag1") String input1,
                   @SideInput("tag2") Integer input2,
-                  BundleFinalizer bundleFinalizer) {}
+                  BundleFinalizer bundleFinalizer,
+                  CausedByDrain causedByDrain) {}
             }.getClass());
 
-    assertThat(sig.processElement().extraParameters().size(), equalTo(9));
+    assertThat(sig.processElement().extraParameters().size(), equalTo(10));
     assertThat(sig.processElement().extraParameters().get(0), instanceOf(ElementParameter.class));
     assertThat(sig.processElement().extraParameters().get(1), instanceOf(TimestampParameter.class));
     assertThat(sig.processElement().extraParameters().get(2), instanceOf(WindowParameter.class));
@@ -146,6 +149,8 @@ public void process(
     assertThat(sig.processElement().extraParameters().get(7), instanceOf(SideInputParameter.class));
     assertThat(
         sig.processElement().extraParameters().get(8), instanceOf(BundleFinalizerParameter.class));
+    assertThat(
+        sig.processElement().extraParameters().get(9), instanceOf(CausedByDrainParameter.class));
   }
 
   @Test
@@ -585,6 +590,31 @@ public void onTimer(BoundedWindow w) {}
         instanceOf(WindowParameter.class));
   }
 
+  @Test
+  public void testCausedByDrainOnTimer() throws Exception {
+    final String timerId = "some-timer-id";
+    final String timerDeclarationId = TimerDeclaration.PREFIX + timerId;
+
+    DoFnSignature sig =
+        DoFnSignatures.getSignature(
+            new DoFn<String, String>() {
+
+              @TimerId(timerId)
+              private final TimerSpec myfield1 = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+              @ProcessElement
+              public void process(ProcessContext c) {}
+
+              @OnTimer(timerId)
+              public void onTimer(CausedByDrain causedByDrain) {}
+            }.getClass());
+
+    assertThat(sig.onTimerMethods().get(timerDeclarationId).extraParameters().size(), equalTo(1));
+    assertThat(
+        sig.onTimerMethods().get(timerDeclarationId).extraParameters().get(0),
+        instanceOf(CausedByDrainParameter.class));
+  }
+
   @Test
   public void testAllParamsOnTimer() throws Exception {
     final String timerId = "some-timer-id";

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

@@ -1804,6 +1804,11 @@ public <T> void outputWindowedValue(
       outputTo(consumer, WindowedValues.of(output, timestamp, windows, paneInfo));
     }
 
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      return currentElement.causedByDrain();
+    }
+
     @Override
     public State state(String stateId, boolean alwaysFetched) {
       StateDeclaration stateDeclaration = doFnSignature.stateDeclarations().get(stateId);
@@ -1946,6 +1951,11 @@ public <T> void outputWindowedValue(
     public CausedByDrain causedByDrain() {
       return currentElement.causedByDrain();
     }
+
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      return currentElement.causedByDrain();
+    }
   }
 
   /** Provides base arguments for a {@link DoFnInvoker} for a non-window observing method. */