private Map<KV<String, String>, KV<Double, Integer>> calculatorOutputAsMap( CallSimilarityCalculatorFactory calculatorFactory) throws Exception { DoFnTester<Variant, KV<KV<String, String>, KV<Double, Integer>>> fnTester = DoFnTester.of(new AlleleSimilarityCalculator(calculatorFactory)); List<KV<KV<String, String>, KV<Double, Integer>>> fnOutput = fnTester.processBundle(variants.toArray(new Variant[] {})); Map<KV<String, String>, KV<Double, Integer>> fnOutputMap = newHashMap(); for (KV<KV<String, String>, KV<Double, Integer>> kv : fnOutput) { fnOutputMap.put(kv.getKey(), kv.getValue()); } return fnOutputMap; }
@Test public void processElementsWithPeeks() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.setCloningBehavior(cloning); tester.startBundle(); tester.processElement(1L); tester.processElement(2L); List<String> peek = tester.peekOutputElements(); assertThat(peek, hasItems("1", "2")); tester.processElement(3L); tester.processElement(4L); peek = tester.peekOutputElements(); assertThat(peek, hasItems("1", "2", "3", "4")); List<String> take = tester.takeOutputElements(); assertThat(take, hasItems("1", "2", "3", "4")); assertTrue(tester.peekOutputElements().isEmpty()); assertTrue(tester.takeOutputElements().isEmpty()); tester.processElement(5L); tester.processElement(6L); peek = tester.peekOutputElements();
/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */ @Deprecated public List<TimestampedValue<OutputT>> peekOutputElementsInWindow(BoundedWindow window) { return peekOutputElementsInWindow(mainOutputTag, window); }
/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */ @Deprecated public List<OutputT> processBundle(Iterable<? extends InputT> inputElements) throws Exception { startBundle(); for (InputT inputElement : inputElements) { processElement(inputElement); } finishBundle(); return takeOutputElements(); }
throws Exception { if (state != State.BUNDLE_STARTED) { startBundle(); createProcessContext( ValueInSingleWindow.of(element, timestamp, window, PaneInfo.NO_FIRING)); fnInvoker.invokeProcessElement( unwrapUserCodeException(e);
@Test public void clonePerBundle() throws Exception { try (DoFnTester<Long, String> tester = DoFnTester.of(new CountBundleCallsFn())) { tester.setCloningBehavior(DoFnTester.CloningBehavior.CLONE_PER_BUNDLE); assertThat(tester.processBundle(1L, 2L, 3L), contains("1/0", "1/0", "1/0")); assertThat(tester.processBundle(4L, 5L), contains("1/0", "1/0")); assertThat(tester.processBundle(6L), contains("1/0")); } }
@Test public void fnWithSideInputExplicit() throws Exception { PCollection<Integer> pCollection = p.apply(Create.of(-2)); final PCollectionView<Integer> value = pCollection.apply(View.<Integer>asSingleton().withDefaultValue(0)); try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new SideInputDoFn(value))) { tester.setSideInput(value, GlobalWindow.INSTANCE, -2); tester.processElement(16); tester.processElement(32); tester.processElement(64); tester.processElement(128); tester.finishBundle(); assertThat(tester.peekOutputElements(), containsInAnyOrder(-2, -2, -2, -2)); } }
@Test public void runQuery() throws Exception { SpannerIO.Read read = SpannerIO.read() .withProjectId("test") .withInstanceId("123") .withDatabaseId("aaa") .withQuery("SELECT * FROM users") .withServiceFactory(serviceFactory); List<Partition> fakePartitions = Arrays.asList(mock(Partition.class), mock(Partition.class), mock(Partition.class)); BatchTransactionId id = mock(BatchTransactionId.class); Transaction tx = Transaction.create(id); PCollectionView<Transaction> txView = pipeline.apply(Create.of(tx)).apply(View.<Transaction>asSingleton()); BatchSpannerRead.GeneratePartitionsFn fn = new BatchSpannerRead.GeneratePartitionsFn(read.getSpannerConfig(), txView); DoFnTester<ReadOperation, Partition> fnTester = DoFnTester.of(fn); fnTester.setSideInput(txView, GlobalWindow.INSTANCE, tx); when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(id)).thenReturn(mockBatchTx); when(mockBatchTx.partitionQuery(any(PartitionOptions.class), any(Statement.class))) .thenReturn(fakePartitions); List<Partition> result = fnTester.processBundle(read.getReadOperation()); assertThat(result, Matchers.containsInAnyOrder(fakePartitions.toArray())); verify(serviceFactory.mockBatchClient()).batchReadOnlyTransaction(id); verify(mockBatchTx) .partitionQuery(any(PartitionOptions.class), eq(Statement.of("SELECT * " + "FROM users"))); }
@Test public void processBundle() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.setCloningBehavior(cloning); // processBundle() returns all the output like takeOutputElements(). assertThat(tester.processBundle(1L, 2L, 3L, 4L), hasItems("1", "2", "3", "4")); // peek now returns nothing. assertTrue(tester.peekOutputElements().isEmpty()); } } }
@Test public void processElement() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.setCloningBehavior(cloning); tester.processElement(1L); List<String> take = tester.takeOutputElements(); assertThat(take, hasItems("1")); // Following takeOutputElements(), neither takeOutputElements() // nor peekOutputElements() return anything. assertTrue(tester.takeOutputElements().isEmpty()); assertTrue(tester.peekOutputElements().isEmpty()); } } }
@Test public void processElementWithOutputTimestamp() throws Exception { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.processElement(1L); tester.processElement(2L); List<TimestampedValue<String>> peek = tester.peekOutputElementsWithTimestamp(); TimestampedValue<String> one = TimestampedValue.of("1", new Instant(1000L)); TimestampedValue<String> two = TimestampedValue.of("2", new Instant(2000L)); assertThat(peek, hasItems(one, two)); tester.processElement(3L); tester.processElement(4L); TimestampedValue<String> three = TimestampedValue.of("3", new Instant(3000L)); TimestampedValue<String> four = TimestampedValue.of("4", new Instant(4000L)); peek = tester.peekOutputElementsWithTimestamp(); assertThat(peek, hasItems(one, two, three, four)); List<TimestampedValue<String>> take = tester.takeOutputElementsWithTimestamp(); assertThat(take, hasItems(one, two, three, four)); // Following takeOutputElementsWithTimestamp(), neither takeOutputElementsWithTimestamp() // nor peekOutputElementsWithTimestamp() return anything. assertTrue(tester.takeOutputElementsWithTimestamp().isEmpty()); assertTrue(tester.peekOutputElementsWithTimestamp().isEmpty()); // peekOutputElements() and takeOutputElements() also return nothing. assertTrue(tester.peekOutputElements().isEmpty()); assertTrue(tester.takeOutputElements().isEmpty()); } }
@Test public void peekValuesInWindow() throws Exception { try (DoFnTester<Long, String> tester = DoFnTester.of(new CounterDoFn())) { tester.startBundle(); tester.processElement(1L); tester.processElement(2L); tester.finishBundle(); assertThat( tester.peekOutputElementsInWindow(GlobalWindow.INSTANCE), containsInAnyOrder( TimestampedValue.of("1", new Instant(1000L)), TimestampedValue.of("2", new Instant(2000L)))); assertThat( tester.peekOutputElementsInWindow(new IntervalWindow(new Instant(0L), new Instant(10L))), Matchers.emptyIterable()); } }
try (DoFnTester<String, Void> fnTester = DoFnTester.of(new Write.WriteFn(write))) { List<String> input = ElasticSearchIOTestUtils.createDocuments( long batchInserted = 0; for (String document : input) { fnTester.processElement(document); numDocsProcessed++; sizeProcessed += document.getBytes(StandardCharsets.UTF_8).length;
Iterable<WriteBundlesToFiles.Result<TableDestination>>, KV<ShardedKey<TableDestination>, List<String>>> tester = DoFnTester.of(writePartition); tester.setSideInput(tempFilePrefixView, GlobalWindow.INSTANCE, tempFilePrefix); tester.processElement(files); partitions = tester.takeOutputElements(multiPartitionsTag); } else { partitions = tester.takeOutputElements(singlePartitionTag);
IsmRecord<WindowedValue<Long>>> doFnTester = DoFnTester.of( new BatchViewOverrides.BatchViewAsMultimap.ToIsmRecordForMapLikeDoFn<>( outputForSizeTag, doFnTester.processBundle(inputElements), contains( IsmRecord.of( doFnTester.takeOutputElements(outputForSizeTag), contains( KV.of( doFnTester.takeOutputElements(outputForEntrySetTag), contains( KV.of(
@Test public void fnWithSideInputDefault() throws Exception { PCollection<Integer> pCollection = p.apply(Create.empty(VarIntCoder.of())); final PCollectionView<Integer> value = pCollection.apply(View.<Integer>asSingleton().withDefaultValue(0)); try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new SideInputDoFn(value))) { tester.processElement(1); tester.processElement(2); tester.processElement(4); tester.processElement(8); assertThat(tester.peekOutputElements(), containsInAnyOrder(0, 0, 0, 0)); } }
@Test public void testSupportsWindowParameter() throws Exception { Instant now = Instant.now(); try (DoFnTester<Integer, KV<Integer, BoundedWindow>> tester = DoFnTester.of(new DoFnWithWindowParameter())) { BoundedWindow firstWindow = new IntervalWindow(now, now.plus(Duration.standardMinutes(1))); tester.processWindowedElement(1, now, firstWindow); tester.processWindowedElement(2, now, firstWindow); BoundedWindow secondWindow = new IntervalWindow(now, now.plus(Duration.standardMinutes(4))); tester.processWindowedElement(3, now, secondWindow); tester.finishBundle(); assertThat( tester.peekOutputElementsInWindow(firstWindow), containsInAnyOrder( TimestampedValue.of(KV.of(1, firstWindow), now), TimestampedValue.of(KV.of(2, firstWindow), now))); assertThat( tester.peekOutputElementsInWindow(secondWindow), containsInAnyOrder(TimestampedValue.of(KV.of(3, secondWindow), now))); } }
@Test public void processTimestampedElement() throws Exception { try (DoFnTester<Long, TimestampedValue<Long>> tester = DoFnTester.of(new ReifyTimestamps())) { TimestampedValue<Long> input = TimestampedValue.of(1L, new Instant(100)); tester.processTimestampedElement(input); assertThat(tester.takeOutputElements(), contains(input)); } }
final ProcessFn<InputT, OutputT, RestrictionT, PositionT> processFn = new ProcessFn<>(fn, inputCoder, restrictionCoder, windowingStrategy); this.tester = DoFnTester.of(processFn); this.timerInternals = new InMemoryTimerInternals(); this.stateInternals = new TestInMemoryStateInternals<>("dummy"); new OutputAndTimeBoundedSplittableProcessElementInvoker<>( fn, tester.getPipelineOptions(), new OutputWindowedValueToDoFnTester<>(tester), new SideInputReader() { this.tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE); this.tester.startBundle(); timerInternals.advanceProcessingTime(currentProcessingTime);
/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */ @Deprecated @SafeVarargs public final List<OutputT> processBundle(InputT... inputElements) throws Exception { return processBundle(Arrays.asList(inputElements)); }