Adds tests for the watermark hold (previously untested)

Author: jkff

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

@@ -72,6 +72,16 @@ public K getKey() {
     return key;
   }
 
+  @Nullable
+  public Instant getWatermarkHold() {
+    for (State state : this.inMemoryState.values()) {
+      if (state instanceof WatermarkHoldState) {
+        return ((WatermarkHoldState<?>) state).read();
+      }
+    }
+    return null;
+  }
+
   /**
    * Interface common to all in-memory state cells. Includes ability to see whether a cell has been
    * cleared and the ability to create a clone of the contents.

runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoTest.java

@@ -19,6 +19,7 @@
 
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.hasItem;
+import static org.hamcrest.Matchers.hasItems;
 import static org.hamcrest.Matchers.not;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -36,6 +37,7 @@
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.InstantCoder;
 import org.apache.beam.sdk.coders.SerializableCoder;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
@@ -334,6 +336,9 @@ List<OutputT> takeOutputElements() {
       return tester.takeOutputElements();
     }
 
+    public Instant getWatermarkHold() {
+      return stateInternals.getWatermarkHold();
+    }
   }
 
   private static class OutputWindowedValueToDoFnTester<OutputT>
@@ -429,6 +434,53 @@ public void testTrivialProcessFnPropagatesOutputsWindowsAndTimestamp() throws Ex
     }
   }
 
+  private static class WatermarkUpdateFn extends DoFn<Instant, String> {
+    @ProcessElement
+    public void process(ProcessContext c, OffsetRangeTracker tracker) {
+      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
+        c.updateWatermark(c.element().plus(Duration.standardSeconds(i)));
+        c.output(String.valueOf(i));
+      }
+    }
+
+    @GetInitialRestriction
+    public OffsetRange getInitialRestriction(Instant elem) {
+      throw new IllegalStateException("Expected to be supplied explicitly in this test");
+    }
+
+    @NewTracker
+    public OffsetRangeTracker newTracker(OffsetRange range) {
+      return new OffsetRangeTracker(range);
+    }
+  }
+
+  @Test
+  public void testUpdatesWatermark() throws Exception {
+    DoFn<Instant, String> fn = new WatermarkUpdateFn();
+    Instant base = Instant.now();
+
+    ProcessFnTester<Instant, String, OffsetRange, OffsetRangeTracker> tester =
+        new ProcessFnTester<>(
+            base,
+            fn,
+            InstantCoder.of(),
+            SerializableCoder.of(OffsetRange.class),
+            3,
+            MAX_BUNDLE_DURATION);
+
+    tester.startElement(base, new OffsetRange(0, 8));
+    assertThat(tester.takeOutputElements(), hasItems("0", "1", "2"));
+    assertEquals(base.plus(Duration.standardSeconds(2)), tester.getWatermarkHold());
+
+    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+    assertThat(tester.takeOutputElements(), hasItems("3", "4", "5"));
+    assertEquals(base.plus(Duration.standardSeconds(5)), tester.getWatermarkHold());
+
+    assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+    assertThat(tester.takeOutputElements(), hasItems("6", "7"));
+    assertEquals(null, tester.getWatermarkHold());
+  }
+
   /**
    * A splittable {@link DoFn} that generates the sequence [init, init + total).
    */