Java Code Examples for org.apache.beam.sdk.util.WindowedValue#valueInGlobalWindow()

@Test
public void mutatedDuringProcessElementThrows() {
  WindowedValue<byte[]> element = WindowedValue.valueInGlobalWindow("bar".getBytes(UTF_8));
  CommittedBundle<byte[]> elements =
      bundleFactory.createBundle(pcollection).add(element).commit(Instant.now());

  ModelEnforcement<byte[]> enforcement = factory.forBundle(elements, consumer);
  enforcement.beforeElement(element);
  element.getValue()[0] = 'f';
  thrown.expect(IllegalMutationException.class);
  thrown.expectMessage(consumer.getFullName());
  thrown.expectMessage("illegaly mutated");
  thrown.expectMessage("Input values must not be mutated");
  enforcement.afterElement(element);
  enforcement.afterFinish(
      elements,
      StepTransformResult.<byte[]>withoutHold(consumer).build(),
      Collections.emptyList());
}
 

@Test
public void expectedInputsAreSent() throws Exception {
  SparkExecutableStageFunction<Integer, ?> function = getFunction(Collections.emptyMap());

  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle);

  @SuppressWarnings("unchecked")
  FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
  when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of(inputId, receiver));

  WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
  WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
  WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
  function.call(Arrays.asList(one, two, three).iterator());

  verify(receiver).accept(one);
  verify(receiver).accept(two);
  verify(receiver).accept(three);
  verifyNoMoreInteractions(receiver);
}
 

@Test
public void expectedInputsAreSent() throws Exception {
  FlinkExecutableStageFunction<Integer> function = getFunction(Collections.emptyMap());
  function.open(new Configuration());

  @SuppressWarnings("unchecked")
  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);

  @SuppressWarnings("unchecked")
  FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
  when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));

  WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
  WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
  WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
  function.mapPartition(Arrays.asList(one, two, three), collector);

  verify(receiver).accept(one);
  verify(receiver).accept(two);
  verify(receiver).accept(three);
  verifyNoMoreInteractions(receiver);
}
 

@Override
public WindowedValue<T> readCurrent() {
  if (finished) {
    throw new IllegalStateException("Bounded reader read all elements");
  }

  final T elem = reader.getCurrent();

  try {
    finished = !reader.advance();
  } catch (final IOException e) {
    e.printStackTrace();
    throw new RuntimeException(e);
  }

  return WindowedValue.valueInGlobalWindow(elem);
}
 

@Test
public void delegatesToUnderlying() throws Exception {
  ParDoEvaluator<Object> underlying = mock(ParDoEvaluator.class);
  DoFn<?, ?> original = lifecycleManager.get();
  TransformEvaluator<Object> evaluator =
      DoFnLifecycleManagerRemovingTransformEvaluator.wrapping(underlying, lifecycleManager);
  WindowedValue<Object> first = WindowedValue.valueInGlobalWindow(new Object());
  WindowedValue<Object> second = WindowedValue.valueInGlobalWindow(new Object());

  evaluator.processElement(first);
  verify(underlying).processElement(first);

  evaluator.processElement(second);
  verify(underlying).processElement(second);

  evaluator.finishBundle();
  verify(underlying).finishBundle();
}
 

@SuppressWarnings("unchecked")
@Test
public void withElementsShouldReturnIndependentBundle() {
  WindowedValue<Integer> firstValue = WindowedValue.valueInGlobalWindow(1);
  WindowedValue<Integer> secondValue =
      WindowedValue.timestampedValueInGlobalWindow(2, new Instant(1000L));

  CommittedBundle<Integer> committed =
      afterCommitGetElementsShouldHaveAddedElements(ImmutableList.of(firstValue, secondValue));

  WindowedValue<Integer> firstReplacement =
      WindowedValue.of(
          9,
          new Instant(2048L),
          new IntervalWindow(new Instant(2044L), Instant.now()),
          PaneInfo.NO_FIRING);
  WindowedValue<Integer> secondReplacement =
      WindowedValue.timestampedValueInGlobalWindow(-1, Instant.now());
  CommittedBundle<Integer> withed =
      committed.withElements(ImmutableList.of(firstReplacement, secondReplacement));

  assertThat(withed.getElements(), containsInAnyOrder(firstReplacement, secondReplacement));
  assertThat(committed.getElements(), containsInAnyOrder(firstValue, secondValue));
  assertThat(withed.getKey(), equalTo(committed.getKey()));
  assertThat(withed.getPCollection(), equalTo(committed.getPCollection()));
  assertThat(
      withed.getSynchronizedProcessingOutputWatermark(),
      equalTo(committed.getSynchronizedProcessingOutputWatermark()));
  assertThat(withed.getMinimumTimestamp(), equalTo(new Instant(2048L)));
}
 

@Test
public void rootBundleSucceedsIgnoresCoder() {
  WindowedValue<Record> one = WindowedValue.valueInGlobalWindow(new Record());
  WindowedValue<Record> two = WindowedValue.valueInGlobalWindow(new Record());
  CommittedBundle<Record> root =
      factory.<Record>createRootBundle().add(one).add(two).commit(Instant.now());

  assertThat(root.getElements(), containsInAnyOrder(one, two));
}
 

@Override
public WindowedValue<byte[]> nextRecord(WindowedValue<byte[]> windowedValue) {
  checkState(availableOutput);
  availableOutput = false;
  if (windowedValue != null) {
    return windowedValue;
  }
  return WindowedValue.valueInGlobalWindow(new byte[0]);
}
 

@Test
public void expectedInputsAreSent() throws Exception {
  TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
  DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
  ExecutableStageDoFnOperator<Integer, Integer> operator =
      getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

  @SuppressWarnings("unchecked")
  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle);

  @SuppressWarnings("unchecked")
  FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
  when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));

  WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
  WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
  WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);

  OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
      new OneInputStreamOperatorTestHarness<>(operator);

  testHarness.open();

  testHarness.processElement(new StreamRecord<>(one));
  testHarness.processElement(new StreamRecord<>(two));
  testHarness.processElement(new StreamRecord<>(three));

  verify(receiver).accept(one);
  verify(receiver).accept(two);
  verify(receiver).accept(three);
  verifyNoMoreInteractions(receiver);

  testHarness.close();
}
 

@Test
public void updateWatermarkWithUnprocessedElements() {
  WindowedValue<Integer> first = WindowedValue.valueInGlobalWindow(1);
  WindowedValue<Integer> second =
      WindowedValue.timestampedValueInGlobalWindow(2, new Instant(-1000L));
  WindowedValue<Integer> third =
      WindowedValue.timestampedValueInGlobalWindow(3, new Instant(1234L));
  CommittedBundle<Integer> createdBundle =
      bundleFactory
          .createBundle(createdInts)
          .add(first)
          .add(second)
          .add(third)
          .commit(clock.now());

  manager.updateWatermarks(
      null,
      TimerUpdate.empty(),
      graph.getProducer(createdInts),
      null,
      Collections.<CommittedBundle<?>>singleton(createdBundle),
      BoundedWindow.TIMESTAMP_MAX_VALUE);

  CommittedBundle<KV<String, Integer>> keyBundle =
      timestampedBundle(
          keyed, TimestampedValue.of(KV.of("MyKey", 1), BoundedWindow.TIMESTAMP_MIN_VALUE));
  manager.updateWatermarks(
      createdBundle,
      TimerUpdate.empty(),
      graph.getProducer(keyed),
      createdBundle.withElements(ImmutableList.of(second, third)),
      Collections.<CommittedBundle<?>>singleton(keyBundle),
      BoundedWindow.TIMESTAMP_MAX_VALUE);
  TransformWatermarks keyedWatermarks = manager.getWatermarks(graph.getProducer(keyed));
  // the unprocessed second and third are readded to pending
  assertThat(keyedWatermarks.getInputWatermark(), not(greaterThan(new Instant(-1000L))));
}
 

@Test
public void unchangedSucceeds() {
  WindowedValue<byte[]> element = WindowedValue.valueInGlobalWindow("bar".getBytes(UTF_8));
  CommittedBundle<byte[]> elements =
      bundleFactory.createBundle(pcollection).add(element).commit(Instant.now());

  ModelEnforcement<byte[]> enforcement = factory.forBundle(elements, consumer);
  enforcement.beforeElement(element);
  enforcement.afterElement(element);
  enforcement.afterFinish(
      elements,
      StepTransformResult.<byte[]>withoutHold(consumer).build(),
      Collections.emptyList());
}
 

@Test
public void testCombineValuesFnAdd() throws Exception {
  TestReceiver receiver = new TestReceiver();
  MeanInts mean = new MeanInts();

  Combine.CombineFn<Integer, CountSum, String> combiner = mean;

  ParDoFn combineParDoFn =
      createCombineValuesFn(
          CombinePhase.ADD,
          combiner,
          StringUtf8Coder.of(),
          BigEndianIntegerCoder.of(),
          new CountSumCoder(),
          WindowingStrategy.globalDefault());

  combineParDoFn.startBundle(receiver);
  combineParDoFn.processElement(
      WindowedValue.valueInGlobalWindow(KV.of("a", Arrays.asList(5, 6, 7))));
  combineParDoFn.processElement(
      WindowedValue.valueInGlobalWindow(KV.of("b", Arrays.asList(1, 3, 7))));
  combineParDoFn.processElement(
      WindowedValue.valueInGlobalWindow(KV.of("c", Arrays.asList(3, 6, 8, 9))));
  combineParDoFn.finishBundle();

  Object[] expectedReceivedElems = {
    WindowedValue.valueInGlobalWindow(KV.of("a", new CountSum(3, 18))),
    WindowedValue.valueInGlobalWindow(KV.of("b", new CountSum(3, 11))),
    WindowedValue.valueInGlobalWindow(KV.of("c", new CountSum(4, 26)))
  };
  assertArrayEquals(expectedReceivedElems, receiver.receivedElems.toArray());
}
 

@Test
public void testCombineValuesFnMerge() throws Exception {
  TestReceiver receiver = new TestReceiver();
  MeanInts mean = new MeanInts();

  Combine.CombineFn<Integer, CountSum, String> combiner = mean;

  ParDoFn combineParDoFn =
      createCombineValuesFn(
          CombinePhase.MERGE,
          combiner,
          StringUtf8Coder.of(),
          BigEndianIntegerCoder.of(),
          new CountSumCoder(),
          WindowingStrategy.globalDefault());

  combineParDoFn.startBundle(receiver);
  combineParDoFn.processElement(
      WindowedValue.valueInGlobalWindow(
          KV.of(
              "a", Arrays.asList(new CountSum(3, 6), new CountSum(2, 9), new CountSum(1, 12)))));
  combineParDoFn.processElement(
      WindowedValue.valueInGlobalWindow(
          KV.of("b", Arrays.asList(new CountSum(2, 20), new CountSum(1, 1)))));
  combineParDoFn.finishBundle();

  Object[] expectedReceivedElems = {
    WindowedValue.valueInGlobalWindow(KV.of("a", new CountSum(6, 27))),
    WindowedValue.valueInGlobalWindow(KV.of("b", new CountSum(3, 21))),
  };
  assertArrayEquals(expectedReceivedElems, receiver.receivedElems.toArray());
}
 

@Test
public void getElementsAfterAddShouldReturnAddedElements() {
  WindowedValue<Integer> firstValue = WindowedValue.valueInGlobalWindow(1);
  WindowedValue<Integer> secondValue =
      WindowedValue.timestampedValueInGlobalWindow(2, new Instant(1000L));

  afterCommitGetElementsShouldHaveAddedElements(ImmutableList.of(firstValue, secondValue));
}
 

@Override
protected WindowedValue<Long> decodeMessage(Windmill.Message message) {
  return WindowedValue.valueInGlobalWindow(message.getTimestamp());
}
 

@Test
public void evaluatorClosesReaderAndResumesFromCheckpoint() throws Exception {
  ContiguousSet<Long> elems = ContiguousSet.create(Range.closed(0L, 20L), DiscreteDomain.longs());
  TestUnboundedSource<Long> source =
      new TestUnboundedSource<>(BigEndianLongCoder.of(), elems.toArray(new Long[0]));

  PCollection<Long> pcollection = p.apply(Read.from(source));
  AppliedPTransform<?, ?, ?> sourceTransform = DirectGraphs.getGraph(p).getProducer(pcollection);

  when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle());
  UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection);
  when(context.createBundle(pcollection)).thenReturn(output);

  WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>> shard =
      WindowedValue.valueInGlobalWindow(
          UnboundedSourceShard.unstarted(source, NeverDeduplicator.create()));
  CommittedBundle<UnboundedSourceShard<Long, TestCheckpointMark>> inputBundle =
      bundleFactory
          .<UnboundedSourceShard<Long, TestCheckpointMark>>createRootBundle()
          .add(shard)
          .commit(Instant.now());
  UnboundedReadEvaluatorFactory factory =
      new UnboundedReadEvaluatorFactory(context, options, 0.0 /* never reuse */);
  TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> evaluator =
      factory.forApplication(sourceTransform, inputBundle);
  evaluator.processElement(shard);
  TransformResult<UnboundedSourceShard<Long, TestCheckpointMark>> result =
      evaluator.finishBundle();

  CommittedBundle<UnboundedSourceShard<Long, TestCheckpointMark>> residual =
      inputBundle.withElements(
          (Iterable<WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>>)
              result.getUnprocessedElements());

  TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> secondEvaluator =
      factory.forApplication(sourceTransform, residual);
  secondEvaluator.processElement(Iterables.getOnlyElement(residual.getElements()));
  secondEvaluator.finishBundle();

  assertThat(TestUnboundedSource.readerClosedCount, equalTo(2));
  assertThat(
      Iterables.getOnlyElement(residual.getElements()).getValue().getCheckpoint().isFinalized(),
      is(true));
}
 

@Test
public void testOutputReceivers() throws Exception {
  TestDoFn fn =
      new TestDoFn(
          ImmutableList.of(
              new TupleTag<>("tag1"), new TupleTag<>("tag2"), new TupleTag<>("tag3")));
  DoFnInfo<?, ?> fnInfo =
      DoFnInfo.forFn(
          fn,
          WindowingStrategy.globalDefault(),
          null /* side input views */,
          null /* input coder */,
          MAIN_OUTPUT,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());
  TestReceiver receiver = new TestReceiver();
  TestReceiver receiver1 = new TestReceiver();
  TestReceiver receiver2 = new TestReceiver();
  TestReceiver receiver3 = new TestReceiver();

  ParDoFn userParDoFn =
      new SimpleParDoFn<>(
          options,
          DoFnInstanceManagers.cloningPool(fnInfo),
          new EmptySideInputReader(),
          MAIN_OUTPUT,
          ImmutableMap.of(
              MAIN_OUTPUT,
              0,
              new TupleTag<String>("tag1"),
              1,
              new TupleTag<String>("tag2"),
              2,
              new TupleTag<String>("tag3"),
              3),
          BatchModeExecutionContext.forTesting(options, "testStage")
              .getStepContext(operationContext),
          operationContext,
          DoFnSchemaInformation.create(),
          Collections.emptyMap(),
          SimpleDoFnRunnerFactory.INSTANCE);

  userParDoFn.startBundle(receiver, receiver1, receiver2, receiver3);

  userParDoFn.processElement(WindowedValue.valueInGlobalWindow(3));
  userParDoFn.processElement(WindowedValue.valueInGlobalWindow(42));
  userParDoFn.processElement(WindowedValue.valueInGlobalWindow(666));

  userParDoFn.finishBundle();

  Object[] expectedReceivedElems = {
    WindowedValue.valueInGlobalWindow("processing: 3"),
    WindowedValue.valueInGlobalWindow("processing: 42"),
    WindowedValue.valueInGlobalWindow("processing: 666"),
    WindowedValue.valueInGlobalWindow("finished"),
  };
  assertArrayEquals(expectedReceivedElems, receiver.receivedElems.toArray());

  Object[] expectedReceivedElems1 = {
    WindowedValue.valueInGlobalWindow("tag1: processing: 3"),
    WindowedValue.valueInGlobalWindow("tag1: processing: 42"),
    WindowedValue.valueInGlobalWindow("tag1: processing: 666"),
    WindowedValue.valueInGlobalWindow("tag1: finished"),
  };
  assertArrayEquals(expectedReceivedElems1, receiver1.receivedElems.toArray());

  Object[] expectedReceivedElems2 = {
    WindowedValue.valueInGlobalWindow("tag2: processing: 3"),
    WindowedValue.valueInGlobalWindow("tag2: processing: 42"),
    WindowedValue.valueInGlobalWindow("tag2: processing: 666"),
    WindowedValue.valueInGlobalWindow("tag2: finished"),
  };
  assertArrayEquals(expectedReceivedElems2, receiver2.receivedElems.toArray());

  Object[] expectedReceivedElems3 = {
    WindowedValue.valueInGlobalWindow("tag3: processing: 3"),
    WindowedValue.valueInGlobalWindow("tag3: processing: 42"),
    WindowedValue.valueInGlobalWindow("tag3: processing: 666"),
    WindowedValue.valueInGlobalWindow("tag3: finished"),
  };
  assertArrayEquals(expectedReceivedElems3, receiver3.receivedElems.toArray());
}
 

@Test
public void outputsAreTaggedCorrectly() throws Exception {
  WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
  WindowedValue<Integer> four = WindowedValue.valueInGlobalWindow(4);
  WindowedValue<Integer> five = WindowedValue.valueInGlobalWindow(5);
  Map<String, Integer> outputTagMap =
      ImmutableMap.of(
          "one", 1,
          "two", 2,
          "three", 3);

  // We use a real StageBundleFactory here in order to exercise the output receiver factory.
  StageBundleFactory stageBundleFactory =
      new StageBundleFactory() {

        private boolean once;

        @Override
        public RemoteBundle getBundle(
            OutputReceiverFactory receiverFactory,
            TimerReceiverFactory timerReceiverFactory,
            StateRequestHandler stateRequestHandler,
            BundleProgressHandler progressHandler) {
          return new RemoteBundle() {
            @Override
            public String getId() {
              return "bundle-id";
            }

            @Override
            public Map<String, FnDataReceiver> getInputReceivers() {
              return ImmutableMap.of(
                  "input",
                  input -> {
                    /* Ignore input*/
                  });
            }

            @Override
            public Map<KV<String, String>, FnDataReceiver<Timer>> getTimerReceivers() {
              return Collections.emptyMap();
            }

            @Override
            public void requestProgress() {
              throw new UnsupportedOperationException();
            }

            @Override
            public void split(double fractionOfRemainder) {
              throw new UnsupportedOperationException();
            }

            @Override
            public void close() throws Exception {
              if (once) {
                return;
              }
              // Emit all values to the runner when the bundle is closed.
              receiverFactory.create("one").accept(three);
              receiverFactory.create("two").accept(four);
              receiverFactory.create("three").accept(five);
              once = true;
            }
          };
        }

        @Override
        public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor
            getProcessBundleDescriptor() {
          return processBundleDescriptor;
        }

        @Override
        public void close() throws Exception {}
      };
  // Wire the stage bundle factory into our context.
  when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);

  FlinkExecutableStageFunction<Integer> function = getFunction(outputTagMap);
  function.open(new Configuration());

  if (isStateful) {
    function.reduce(Collections.emptyList(), collector);
  } else {
    function.mapPartition(Collections.emptyList(), collector);
  }
  // Ensure that the tagged values sent to the collector have the correct union tags as specified
  // in the output map.
  verify(collector).collect(new RawUnionValue(1, three));
  verify(collector).collect(new RawUnionValue(2, four));
  verify(collector).collect(new RawUnionValue(3, five));
  verifyNoMoreInteractions(collector);
}
 

@Test
public void outputsAreTaggedCorrectly() throws Exception {
  WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
  WindowedValue<Integer> four = WindowedValue.valueInGlobalWindow(4);
  WindowedValue<Integer> five = WindowedValue.valueInGlobalWindow(5);
  Map<String, Integer> outputTagMap =
      ImmutableMap.of(
          "one", 1,
          "two", 2,
          "three", 3);

  // We use a real StageBundleFactory here in order to exercise the output receiver factory.
  StageBundleFactory stageBundleFactory =
      new StageBundleFactory() {

        private boolean once;

        @Override
        public RemoteBundle getBundle(
            OutputReceiverFactory receiverFactory,
            TimerReceiverFactory timerReceiverFactory,
            StateRequestHandler stateRequestHandler,
            BundleProgressHandler progressHandler) {
          return new RemoteBundle() {
            @Override
            public String getId() {
              return "bundle-id";
            }

            @Override
            public Map<String, FnDataReceiver> getInputReceivers() {
              return ImmutableMap.of(
                  "input",
                  input -> {
                    /* Ignore input*/
                  });
            }

            @Override
            public Map<KV<String, String>, FnDataReceiver<Timer>> getTimerReceivers() {
              return Collections.emptyMap();
            }

            @Override
            public void requestProgress() {
              throw new UnsupportedOperationException();
            }

            @Override
            public void split(double fractionOfRemainder) {
              throw new UnsupportedOperationException();
            }

            @Override
            public void close() throws Exception {
              if (once) {
                return;
              }
              // Emit all values to the runner when the bundle is closed.
              receiverFactory.create("one").accept(three);
              receiverFactory.create("two").accept(four);
              receiverFactory.create("three").accept(five);
              once = true;
            }
          };
        }

        @Override
        public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor
            getProcessBundleDescriptor() {
          return Mockito.mock(ProcessBundleDescriptors.ExecutableProcessBundleDescriptor.class);
        }

        @Override
        public void close() {}
      };
  when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);

  SparkExecutableStageFunction<Integer, ?> function = getFunction(outputTagMap);
  Iterator<RawUnionValue> iterator = function.call(Collections.emptyIterator());
  Iterable<RawUnionValue> iterable = () -> iterator;

  assertThat(
      iterable,
      contains(
          new RawUnionValue(1, three), new RawUnionValue(2, four), new RawUnionValue(3, five)));
}
 

@Test
public void testSideInputs() {
  // mock context
  final Transform.Context context = mock(Transform.Context.class);
  TupleTag<Tuple<String, Iterable<String>>> outputTag = new TupleTag<>("main-output");

  WindowedValue<String> firstElement = WindowedValue.valueInGlobalWindow("first");
  WindowedValue<String> secondElement = WindowedValue.valueInGlobalWindow("second");

  SideInputElement firstSideinput = new SideInputElement<>(0, ImmutableList.of("1"));
  SideInputElement secondSideinput = new SideInputElement(1, ImmutableList.of("2"));

  final Map<Integer, PCollectionView<?>> sideInputMap = new HashMap<>();
  sideInputMap.put(firstSideinput.getSideInputIndex(), view1);
  sideInputMap.put(secondSideinput.getSideInputIndex(), view2);
  final PushBackDoFnTransform<String, String> doFnTransform =
    new PushBackDoFnTransform(
      new SimpleSideInputDoFn<String>(view1, view2),
      NULL_INPUT_CODER,
      NULL_OUTPUT_CODERS,
      outputTag,
      Collections.emptyList(),
      WindowingStrategy.globalDefault(),
      sideInputMap, /* side inputs */
      PipelineOptionsFactory.as(NemoPipelineOptions.class),
      DisplayData.none(),
      DoFnSchemaInformation.create(),
      Collections.emptyMap());

  final TestOutputCollector<String> oc = new TestOutputCollector<>();
  doFnTransform.prepare(context, oc);

  // Main input first, Side inputs later
  doFnTransform.onData(firstElement);

  doFnTransform.onData(WindowedValue.valueInGlobalWindow(firstSideinput));
  doFnTransform.onData(WindowedValue.valueInGlobalWindow(secondSideinput));
  assertEquals(
    WindowedValue.valueInGlobalWindow(
      concat(firstElement.getValue(), firstSideinput.getSideInputValue(), secondSideinput.getSideInputValue())),
    oc.getOutput().get(0));

  // Side inputs first, Main input later
  doFnTransform.onData(secondElement);
  assertEquals(
    WindowedValue.valueInGlobalWindow(
      concat(secondElement.getValue(), firstSideinput.getSideInputValue(), secondSideinput.getSideInputValue())),
    oc.getOutput().get(1));

  // There should be only 2 final outputs
  assertEquals(2, oc.getOutput().size());

  // The side inputs should be "READY"
  assertTrue(doFnTransform.getSideInputReader().isReady(view1, GlobalWindow.INSTANCE));
  assertTrue(doFnTransform.getSideInputReader().isReady(view2, GlobalWindow.INSTANCE));

  // This watermark should remove the side inputs. (Now should be "NOT READY")
  doFnTransform.onWatermark(new Watermark(GlobalWindow.TIMESTAMP_MAX_VALUE.getMillis()));
  Iterable materializedSideInput1 = doFnTransform.getSideInputReader().get(view1, GlobalWindow.INSTANCE);
  Iterable materializedSideInput2 = doFnTransform.getSideInputReader().get(view2, GlobalWindow.INSTANCE);
  assertFalse(materializedSideInput1.iterator().hasNext());
  assertFalse(materializedSideInput2.iterator().hasNext());

  // There should be only 2 final outputs
  doFnTransform.close();
  assertEquals(2, oc.getOutput().size());
}