content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.beam.sdk.transforms;
  
  import static org.apache.beam.sdk.transforms.Requirements.requiresSideInputs;
  import static org.apache.beam.sdk.transforms.Watch.Growth.afterTimeSinceNewOutput;
  import static org.apache.beam.sdk.transforms.Watch.Growth.allOf;
  import static org.apache.beam.sdk.transforms.Watch.Growth.eitherOf;
  import static org.apache.beam.sdk.transforms.Watch.Growth.never;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.containsInAnyOrder;
  import static org.joda.time.Duration.standardSeconds;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertNull;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  import static org.mockito.Mockito.mock;
  
  import java.io.IOException;
  import java.io.Serializable;
  import java.util.Arrays;
  import java.util.List;
  import java.util.UUID;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.coders.KvCoder;
  import org.apache.beam.sdk.coders.StringUtf8Coder;
  import org.apache.beam.sdk.coders.VarIntCoder;
  import org.apache.beam.sdk.testing.CoderProperties;
  import org.apache.beam.sdk.testing.NeedsRunner;
  import org.apache.beam.sdk.testing.PAssert;
  import org.apache.beam.sdk.testing.TestPipeline;
  import org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo;
  import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation;
  import org.apache.beam.sdk.transforms.Watch.Growth;
  import org.apache.beam.sdk.transforms.Watch.Growth.PollFn;
  import org.apache.beam.sdk.transforms.Watch.Growth.PollResult;
  import org.apache.beam.sdk.transforms.Watch.GrowthState;
  import org.apache.beam.sdk.transforms.Watch.GrowthTracker;
  import org.apache.beam.sdk.transforms.Watch.NonPollingGrowthState;
  import org.apache.beam.sdk.transforms.Watch.PollingGrowthState;
  import org.apache.beam.sdk.transforms.Watch.WatchGrowthFn;
  import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
  import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.PCollectionView;
  import org.apache.beam.sdk.values.TimestampedValue;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Funnel;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Funnels;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.HashCode;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.joda.time.ReadableDuration;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.experimental.categories.Category;
  import org.junit.runner.RunWith;
  import org.junit.runners.JUnit4;
  
  /** Tests for {@link Watch}. */
  @RunWith(JUnit4.class)
  public class WatchTest implements Serializable {
  
    @Rule public transient TestPipeline p = TestPipeline.create();
  
    @Test
    @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
    public void testSinglePollMultipleInputs() {
      PCollection<KV<String, String>> res =
          p.apply(Create.of("a", "b"))
              .apply(
                  Watch.growthOf(
                          new PollFn<String, String>() {
                            @Override
                            public PollResult<String> apply(String element, Context c)
                                throws Exception {
                              return PollResult.complete(
                                  Instant.now(), Arrays.asList(element + ".foo", element + ".bar"));
                            }
                          })
                      .withPollInterval(Duration.ZERO));
  
      PAssert.that(res)
          .containsInAnyOrder(
              Arrays.asList(
                  KV.of("a", "a.foo"), KV.of("a", "a.bar"),
                  KV.of("b", "b.foo"), KV.of("b", "b.bar")));
  
      p.run();
    }
  
    @Test
    @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
    public void testSinglePollMultipleInputsWithSideInput() {
      final PCollectionView<String> moo =
          p.apply("moo", Create.of("moo")).apply("moo singleton", View.asSingleton());
      final PCollectionView<String> zoo =
          p.apply("zoo", Create.of("zoo")).apply("zoo singleton", View.asSingleton());
      PCollection<KV<String, String>> res =
          p.apply("input", Create.of("a", "b"))
              .apply(
                  Watch.growthOf(
                          new PollFn<String, String>() {
                            @Override
                            public PollResult<String> apply(String element, Context c)
                                throws Exception {
                              return PollResult.complete(
                                  Instant.now(),
                                  Arrays.asList(
                                      element + " " + c.sideInput(moo) + " " + c.sideInput(zoo)));
                            }
                          },
                          requiresSideInputs(moo, zoo))
                      .withPollInterval(Duration.ZERO));
  
      PAssert.that(res)
          .containsInAnyOrder(Arrays.asList(KV.of("a", "a moo zoo"), KV.of("b", "b moo zoo")));
  
      p.run();
    }
  
    @Test
    @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
    public void testMultiplePollsWithTerminationBecauseOutputIsFinal() {
      testMultiplePolls(false);
    }
  
    @Test
    @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
    public void testMultiplePollsWithTerminationDueToTerminationCondition() {
      testMultiplePolls(true);
    }
  
    private void testMultiplePolls(boolean terminationConditionElapsesBeforeOutputIsFinal) {
      List<Integer> all = Arrays.asList(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
  
      PCollection<Integer> res =
          p.apply(Create.of("a"))
              .apply(
                  Watch.growthOf(
                          new TimedPollFn<String, Integer>(
                              all,
                              standardSeconds(1) /* timeToOutputEverything */,
                              standardSeconds(3) /* timeToDeclareOutputFinal */,
                              standardSeconds(30) /* timeToFail */))
                      .withTerminationPerInput(
                          Growth.afterTotalOf(
                              standardSeconds(
                                  // At 2 seconds, all output has been yielded, but not yet
                                  // declared final - so polling should terminate per termination
                                  // condition.
                                  // At 3 seconds, all output has been yielded (and declared final),
                                  // so polling should terminate because of that without waiting for
                                  // 100 seconds.
                                  terminationConditionElapsesBeforeOutputIsFinal ? 2 : 100)))
                      .withPollInterval(Duration.millis(300))
                      .withOutputCoder(VarIntCoder.of()))
              .apply("Drop input", Values.create());
  
      PAssert.that(res).containsInAnyOrder(all);
  
      p.run();
    }
  
    @Test
    @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
    public void testMultiplePollsWithKeyExtractor() {
      List<KV<Integer, String>> polls =
          Arrays.asList(
              KV.of(0, "0"),
              KV.of(10, "10"),
              KV.of(20, "20"),
              KV.of(30, "30"),
              KV.of(40, "40"),
              KV.of(40, "40.1"),
              KV.of(20, "20.1"),
              KV.of(50, "50"),
              KV.of(10, "10.1"),
              KV.of(10, "10.2"),
              KV.of(60, "60"),
              KV.of(70, "70"),
              KV.of(60, "60.1"),
              KV.of(80, "80"),
              KV.of(40, "40.2"),
              KV.of(90, "90"),
              KV.of(90, "90.1"));
  
      List<Integer> expected = Arrays.asList(0, 10, 20, 30, 40, 50, 60, 70, 80, 90);
  
      PCollection<Integer> res =
          p.apply(Create.of("a"))
              .apply(
                  Watch.growthOf(
                          Contextful.of(
                              new TimedPollFn<String, KV<Integer, String>>(
                                  polls,
                                  standardSeconds(1) /* timeToOutputEverything */,
                                  standardSeconds(3) /* timeToDeclareOutputFinal */,
                                  standardSeconds(30) /* timeToFail */),
                              Requirements.empty()),
                          KV::getKey)
                      .withTerminationPerInput(Growth.afterTotalOf(standardSeconds(5)))
                      .withPollInterval(Duration.millis(100))
                      .withOutputCoder(KvCoder.of(VarIntCoder.of(), StringUtf8Coder.of()))
                      .withOutputKeyCoder(VarIntCoder.of()))
              .apply("Drop input", Values.create())
              .apply("Drop auxiliary string", Keys.create());
  
      PAssert.that(res).containsInAnyOrder(expected);
  
      p.run();
    }
  
    @Test
    @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
    public void testMultiplePollsStopAfterTimeSinceNewOutput() {
      List<Integer> all = Arrays.asList(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
  
      PCollection<Integer> res =
          p.apply(Create.of("a"))
              .apply(
                  Watch.growthOf(
                          new TimedPollFn<String, Integer>(
                              all,
                              standardSeconds(1) /* timeToOutputEverything */,
                              // Never declare output final
                              standardSeconds(1000) /* timeToDeclareOutputFinal */,
                              standardSeconds(30) /* timeToFail */))
                      // Should terminate after 4 seconds - earlier than timeToFail
                      .withTerminationPerInput(afterTimeSinceNewOutput(standardSeconds(3)))
                      .withPollInterval(Duration.millis(300))
                      .withOutputCoder(VarIntCoder.of()))
              .apply("Drop input", Values.create());
  
      PAssert.that(res).containsInAnyOrder(all);
  
      p.run();
    }
  
    @Test
    @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
    public void testSinglePollWithManyResults() {
      // More than the default 100 elements per checkpoint for direct runner.
      final long numResults = 3000;
      PCollection<KV<String, Integer>> res =
          p.apply(Create.of("a"))
              .apply(
                  Watch.growthOf(
                          new PollFn<String, KV<String, Integer>>() {
                            @Override
                            public PollResult<KV<String, Integer>> apply(String element, Context c)
                                throws Exception {
                              String pollId = UUID.randomUUID().toString();
                              List<KV<String, Integer>> output = Lists.newArrayList();
                              for (int i = 0; i < numResults; ++i) {
                                output.add(KV.of(pollId, i));
                              }
                              return PollResult.complete(Instant.now(), output);
                            }
                          })
                      .withTerminationPerInput(Growth.afterTotalOf(standardSeconds(1)))
                      .withPollInterval(Duration.millis(1))
                      .withOutputCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())))
              .apply("Drop input", Values.create());
  
      PAssert.that("Poll called only once", res.apply(Keys.create()))
          .satisfies(
              pollIds -> {
                assertEquals(1, Sets.newHashSet(pollIds).size());
                return null;
              });
      PAssert.that("Yields all expected results", res.apply("Drop poll id", Values.create()))
          .satisfies(
              input -> {
                assertEquals(
                    "Total number of results mismatches",
                    numResults,
                    Lists.newArrayList(input).size());
                assertEquals("Results are not unique", numResults, Sets.newHashSet(input).size());
                return null;
              });
  
      p.run();
    }
  
    @Test
    public void testCoder() throws Exception {
      GrowthState pollingState =
          PollingGrowthState.of(
              ImmutableMap.of(
                  HashCode.fromString("0123456789abcdef0123456789abcdef"), Instant.now(),
                  HashCode.fromString("01230123012301230123012301230123"), Instant.now()),
              Instant.now(),
              "STATE");
      GrowthState nonPollingState =
          NonPollingGrowthState.of(
              Growth.PollResult.incomplete(Instant.now(), Arrays.asList("A", "B")));
      Coder<GrowthState> coder =
          Watch.GrowthStateCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());
  
      CoderProperties.coderDecodeEncodeEqual(coder, pollingState);
      CoderProperties.coderDecodeEncodeEqual(coder, nonPollingState);
    }
  
    /**
     * Gradually emits all items from the given list, pairing each one with a UUID that identifies the
     * round of polling, so a client can check how many rounds of polling there were.
     */
    private static class TimedPollFn<InputT, OutputT> extends PollFn<InputT, OutputT> {
  
      private final Instant baseTime;
      private final List<OutputT> outputs;
      private final Duration timeToOutputEverything;
      private final Duration timeToDeclareOutputFinal;
      private final Duration timeToFail;
  
      public TimedPollFn(
          List<OutputT> outputs,
          Duration timeToOutputEverything,
          Duration timeToDeclareOutputFinal,
          Duration timeToFail) {
        this.baseTime = Instant.now();
        this.outputs = outputs;
        this.timeToOutputEverything = timeToOutputEverything;
        this.timeToDeclareOutputFinal = timeToDeclareOutputFinal;
        this.timeToFail = timeToFail;
      }
  
      @Override
      public PollResult<OutputT> apply(InputT element, Context c) throws Exception {
        Instant now = Instant.now();
        Duration elapsed = new Duration(baseTime, Instant.now());
        if (elapsed.isLongerThan(timeToFail)) {
          fail(
              String.format(
                  "Poll called %s after base time, which is longer than the threshold of %s",
                  elapsed, timeToFail));
        }
  
        double fractionElapsed = 1.0 * elapsed.getMillis() / timeToOutputEverything.getMillis();
        int numToEmit = (int) Math.min(outputs.size(), fractionElapsed * outputs.size());
        List<TimestampedValue<OutputT>> toEmit = Lists.newArrayList();
        for (int i = 0; i < numToEmit; ++i) {
          toEmit.add(TimestampedValue.of(outputs.get(i), now));
        }
        return elapsed.isLongerThan(timeToDeclareOutputFinal)
            ? PollResult.complete(toEmit)
            : PollResult.incomplete(toEmit).withWatermark(now);
      }
    }
  
    @Test
    public void testTerminationConditionsNever() {
      Watch.Growth.Never<Object> c = never();
      Integer state = c.forNewInput(Instant.now(), null);
      assertFalse(c.canStopPolling(Instant.now(), state));
    }
  
    @Test
    public void testTerminationConditionsAfterTotalOf() {
      Instant now = Instant.now();
      Watch.Growth.AfterTotalOf<Object> c = Growth.afterTotalOf(standardSeconds(5));
      KV<Instant, ReadableDuration> state = c.forNewInput(now, null);
      assertFalse(c.canStopPolling(now, state));
      assertFalse(c.canStopPolling(now.plus(standardSeconds(3)), state));
      assertTrue(c.canStopPolling(now.plus(standardSeconds(6)), state));
    }
  
    @Test
    public void testTerminationConditionsAfterTimeSinceNewOutput() {
      Instant now = Instant.now();
      Watch.Growth.AfterTimeSinceNewOutput<Object> c = afterTimeSinceNewOutput(standardSeconds(5));
      KV<Instant, ReadableDuration> state = c.forNewInput(now, null);
      assertFalse(c.canStopPolling(now, state));
      assertFalse(c.canStopPolling(now.plus(standardSeconds(3)), state));
      assertFalse(c.canStopPolling(now.plus(standardSeconds(6)), state));
  
      state = c.onSeenNewOutput(now.plus(standardSeconds(3)), state);
      assertFalse(c.canStopPolling(now.plus(standardSeconds(3)), state));
      assertFalse(c.canStopPolling(now.plus(standardSeconds(6)), state));
      assertTrue(c.canStopPolling(now.plus(standardSeconds(9)), state));
  
      state = c.onSeenNewOutput(now.plus(standardSeconds(5)), state);
      assertFalse(c.canStopPolling(now.plus(standardSeconds(3)), state));
      assertFalse(c.canStopPolling(now.plus(standardSeconds(6)), state));
      assertFalse(c.canStopPolling(now.plus(standardSeconds(9)), state));
      assertTrue(c.canStopPolling(now.plus(standardSeconds(11)), state));
    }
  
    @Test
    public void testTerminationConditionsEitherOf() {
      Instant now = Instant.now();
      Watch.Growth.AfterTotalOf<Object> a = Growth.afterTotalOf(standardSeconds(5));
      Watch.Growth.AfterTotalOf<Object> b = Growth.afterTotalOf(standardSeconds(10));
  
      Watch.Growth.BinaryCombined<
              Object, KV<Instant, ReadableDuration>, KV<Instant, ReadableDuration>>
          c = eitherOf(a, b);
      KV<KV<Instant, ReadableDuration>, KV<Instant, ReadableDuration>> state =
          c.forNewInput(now, null);
      assertFalse(c.canStopPolling(now.plus(standardSeconds(3)), state));
      assertTrue(c.canStopPolling(now.plus(standardSeconds(7)), state));
      assertTrue(c.canStopPolling(now.plus(standardSeconds(12)), state));
    }
  
    @Test
    public void testTerminationConditionsAllOf() {
      Instant now = Instant.now();
      Watch.Growth.AfterTotalOf<Object> a = Growth.afterTotalOf(standardSeconds(5));
      Watch.Growth.AfterTotalOf<Object> b = Growth.afterTotalOf(standardSeconds(10));
  
      Watch.Growth.BinaryCombined<
              Object, KV<Instant, ReadableDuration>, KV<Instant, ReadableDuration>>
          c = allOf(a, b);
      KV<KV<Instant, ReadableDuration>, KV<Instant, ReadableDuration>> state =
          c.forNewInput(now, null);
      assertFalse(c.canStopPolling(now.plus(standardSeconds(3)), state));
      assertFalse(c.canStopPolling(now.plus(standardSeconds(7)), state));
      assertTrue(c.canStopPolling(now.plus(standardSeconds(12)), state));
    }
  
    private static GrowthTracker<String, Integer> newTracker(GrowthState state) {
      Funnel<String> coderFunnel =
          (from, into) -> {
            try {
              StringUtf8Coder.of().encode(from, Funnels.asOutputStream(into));
            } catch (IOException e) {
              throw new RuntimeException(e);
            }
          };
      return new GrowthTracker<>(state, coderFunnel);
    }
  
    private static HashCode hash128(String value) {
      Funnel<String> coderFunnel =
          (from, into) -> {
            try {
              StringUtf8Coder.of().encode(from, Funnels.asOutputStream(into));
            } catch (IOException e) {
              throw new RuntimeException(e);
            }
          };
      return Hashing.murmur3_128().hashObject(value, coderFunnel);
    }
  
    private static GrowthTracker<String, Integer> newPollingGrowthTracker() {
      return newTracker(PollingGrowthState.of(never().forNewInput(Instant.now(), null)));
    }
  
    @Test
    public void testPollingGrowthTrackerUsesElementTimestampIfNoWatermarkProvided() throws Exception {
      Instant now = Instant.now();
      Watch.Growth<String, String, String> growth =
          Watch.growthOf(
                  new Watch.Growth.PollFn<String, String>() {
  
                    @Override
                    public PollResult<String> apply(String element, Context c) throws Exception {
                      // We specifically test an unsorted list.
                      return PollResult.incomplete(
                          Arrays.asList(
                              TimestampedValue.of("d", now.plus(standardSeconds(4))),
                              TimestampedValue.of("c", now.plus(standardSeconds(3))),
                              TimestampedValue.of("a", now.plus(standardSeconds(1))),
                              TimestampedValue.of("b", now.plus(standardSeconds(2)))));
                    }
                  })
              .withPollInterval(standardSeconds(10));
      WatchGrowthFn<String, String, String, Integer> growthFn =
          new WatchGrowthFn(
              growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of());
      GrowthTracker<String, Integer> tracker = newPollingGrowthTracker();
      DoFn.ProcessContext context = mock(DoFn.ProcessContext.class);
      ManualWatermarkEstimator<Instant> watermarkEstimator =
          new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
      ProcessContinuation processContinuation =
          growthFn.process(context, tracker, watermarkEstimator);
      assertEquals(now.plus(standardSeconds(1)), watermarkEstimator.currentWatermark());
      assertTrue(processContinuation.shouldResume());
    }
  
    @Test
    public void testPollingGrowthTrackerCheckpointNonEmpty() {
      Instant now = Instant.now();
      GrowthTracker<String, Integer> tracker = newPollingGrowthTracker();
  
      PollResult<String> claim =
          PollResult.incomplete(
                  Arrays.asList(
                      TimestampedValue.of("d", now.plus(standardSeconds(4))),
                      TimestampedValue.of("c", now.plus(standardSeconds(3))),
                      TimestampedValue.of("a", now.plus(standardSeconds(1))),
                      TimestampedValue.of("b", now.plus(standardSeconds(2)))))
              .withWatermark(now.plus(standardSeconds(7)));
  
      assertTrue(tracker.tryClaim(KV.of(claim, 1 /* termination state */)));
  
      PollingGrowthState<Integer> residual =
          (PollingGrowthState<Integer>) tracker.trySplit(0).getResidual();
      NonPollingGrowthState<String> primary =
          (NonPollingGrowthState<String>) tracker.currentRestriction();
      tracker.checkDone();
  
      // Verify primary: should contain what the current tracker claimed, and nothing else.
      assertEquals(claim, primary.getPending());
  
      // Verify residual: should contain what the current tracker didn't claim.
      assertEquals(now.plus(standardSeconds(7)), residual.getPollWatermark());
      assertThat(
          residual.getCompleted().keySet(),
          containsInAnyOrder(hash128("a"), hash128("b"), hash128("c"), hash128("d")));
      assertEquals(1, (int) residual.getTerminationState());
    }
  
    @Test
    public void testPollingGrowthTrackerCheckpointEmpty() {
      GrowthTracker<String, Integer> tracker = newPollingGrowthTracker();
  
      PollingGrowthState<Integer> residual =
          (PollingGrowthState<Integer>) tracker.trySplit(0).getResidual();
      GrowthState primary = tracker.currentRestriction();
      tracker.checkDone();
  
      // Verify primary: should contain what the current tracker claimed, and nothing else.
      assertEquals(GrowthTracker.EMPTY_STATE, primary);
  
      // Verify residual: should contain what the current tracker didn't claim.
      assertNull(residual.getPollWatermark());
      assertEquals(0, residual.getCompleted().size());
      assertEquals(0, (int) residual.getTerminationState());
    }
  
    @Test
    public void testPollingGrowthTrackerHashAlreadyClaimed() {
      Instant now = Instant.now();
      GrowthTracker<String, Integer> tracker = newPollingGrowthTracker();
  
      PollResult<String> claim =
          PollResult.incomplete(
                  Arrays.asList(
                      TimestampedValue.of("d", now.plus(standardSeconds(4))),
                      TimestampedValue.of("c", now.plus(standardSeconds(3))),
                      TimestampedValue.of("a", now.plus(standardSeconds(1))),
                      TimestampedValue.of("b", now.plus(standardSeconds(2)))))
              .withWatermark(now.plus(standardSeconds(7)));
  
      assertTrue(tracker.tryClaim(KV.of(claim, 1 /* termination state */)));
  
      PollingGrowthState<Integer> residual =
          (PollingGrowthState<Integer>) tracker.trySplit(0).getResidual();
  
      assertFalse(newTracker(residual).tryClaim(KV.of(claim, 2)));
    }
  
    @Test
    public void testNonPollingGrowthTrackerIgnoresWatermark() throws Exception {
      Instant now = Instant.now();
      PollResult<String> claim =
          PollResult.incomplete(
                  Arrays.asList(
                      TimestampedValue.of("d", now.plus(standardSeconds(4))),
                      TimestampedValue.of("c", now.plus(standardSeconds(3))),
                      TimestampedValue.of("a", now.plus(standardSeconds(1))),
                      TimestampedValue.of("b", now.plus(standardSeconds(2)))))
              .withWatermark(now.plus(standardSeconds(7)));
  
      Watch.Growth<String, String, String> growth =
          Watch.growthOf(
                  new Watch.Growth.PollFn<String, String>() {
  
                    @Override
                    public PollResult<String> apply(String element, Context c) throws Exception {
                      fail("Never expected to be invoked for NonPollingGrowthState.");
                      return null;
                    }
                  })
              .withPollInterval(standardSeconds(10));
      GrowthTracker<String, Integer> tracker = newTracker(NonPollingGrowthState.of(claim));
      WatchGrowthFn<String, String, String, Integer> growthFn =
          new WatchGrowthFn(
              growth, StringUtf8Coder.of(), SerializableFunctions.identity(), StringUtf8Coder.of());
      DoFn.ProcessContext context = mock(DoFn.ProcessContext.class);
      ManualWatermarkEstimator<Instant> watermarkEstimator =
          new WatermarkEstimators.Manual(BoundedWindow.TIMESTAMP_MIN_VALUE);
      ProcessContinuation processContinuation =
          growthFn.process(context, tracker, watermarkEstimator);
      assertEquals(BoundedWindow.TIMESTAMP_MIN_VALUE, watermarkEstimator.currentWatermark());
      assertFalse(processContinuation.shouldResume());
    }
  
    @Test
    public void testNonPollingGrowthTrackerCheckpointNonEmpty() {
      Instant now = Instant.now();
      PollResult<String> claim =
          PollResult.incomplete(
                  Arrays.asList(
                      TimestampedValue.of("d", now.plus(standardSeconds(4))),
                      TimestampedValue.of("c", now.plus(standardSeconds(3))),
                      TimestampedValue.of("a", now.plus(standardSeconds(1))),
                      TimestampedValue.of("b", now.plus(standardSeconds(2)))))
              .withWatermark(now.plus(standardSeconds(7)));
  
      GrowthTracker<String, Integer> tracker = newTracker(NonPollingGrowthState.of(claim));
  
      assertTrue(tracker.tryClaim(KV.of(claim, 1 /* termination state */)));
      GrowthState residual = tracker.trySplit(0).getResidual();
      NonPollingGrowthState<String> primary =
          (NonPollingGrowthState<String>) tracker.currentRestriction();
      tracker.checkDone();
  
      // Verify primary: should contain what the current tracker claimed, and nothing else.
      assertEquals(claim, primary.getPending());
  
      // Verify residual: should contain what the current tracker didn't claim.
      assertEquals(GrowthTracker.EMPTY_STATE, residual);
    }
  
    @Test
    public void testNonPollingGrowthTrackerCheckpointEmpty() {
      Instant now = Instant.now();
      PollResult<String> claim =
          PollResult.incomplete(
                  Arrays.asList(
                      TimestampedValue.of("d", now.plus(standardSeconds(4))),
                      TimestampedValue.of("c", now.plus(standardSeconds(3))),
                      TimestampedValue.of("a", now.plus(standardSeconds(1))),
                      TimestampedValue.of("b", now.plus(standardSeconds(2)))))
              .withWatermark(now.plus(standardSeconds(7)));
  
      GrowthTracker<String, Integer> tracker = newTracker(NonPollingGrowthState.of(claim));
  
      NonPollingGrowthState<String> residual =
          (NonPollingGrowthState<String>) tracker.trySplit(0).getResidual();
      GrowthState primary = tracker.currentRestriction();
      tracker.checkDone();
  
      // Verify primary: should contain what the current tracker claimed, and nothing else.
      assertEquals(GrowthTracker.EMPTY_STATE, primary);
  
      // Verify residual: should contain what the current tracker didn't claim.
      assertEquals(claim, residual.getPending());
    }
  
    @Test
    public void testNonPollingGrowthTrackerFailedToClaimOtherPollResult() {
      Instant now = Instant.now();
      PollResult<String> claim =
          PollResult.incomplete(
                  Arrays.asList(
                      TimestampedValue.of("d", now.plus(standardSeconds(4))),
                      TimestampedValue.of("c", now.plus(standardSeconds(3))),
                      TimestampedValue.of("a", now.plus(standardSeconds(1))),
                      TimestampedValue.of("b", now.plus(standardSeconds(2)))))
              .withWatermark(now.plus(standardSeconds(7)));
  
      GrowthTracker<String, Integer> tracker = newTracker(NonPollingGrowthState.of(claim));
  
      PollResult<String> otherClaim =
          PollResult.incomplete(
                  Arrays.asList(
                      TimestampedValue.of("x", now.plus(standardSeconds(14))),
                      TimestampedValue.of("y", now.plus(standardSeconds(13))),
                      TimestampedValue.of("z", now.plus(standardSeconds(12)))))
              .withWatermark(now.plus(standardSeconds(17)));
      assertFalse(tracker.tryClaim(KV.of(otherClaim, 1)));
    }
  }
  