/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */ @Deprecated @SafeVarargs public final List<OutputT> processBundle(InputT... inputElements) throws Exception { return processBundle(Arrays.asList(inputElements)); }
@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 cloneOnce() throws Exception { try (DoFnTester<Long, String> tester = DoFnTester.of(new CountBundleCallsFn())) { tester.setCloningBehavior(DoFnTester.CloningBehavior.CLONE_ONCE); assertThat(tester.processBundle(1L, 2L, 3L), contains("1/0", "1/0", "1/0")); assertThat(tester.processBundle(4L, 5L), contains("2/1", "2/1")); assertThat(tester.processBundle(6L), contains("3/2")); } }
@Test public void testSupportsFinishBundleOutput() throws Exception { for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) { try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new BundleCounterDoFn())) { tester.setCloningBehavior(cloning); assertThat(tester.processBundle(1, 2, 3, 4), contains(4)); assertThat(tester.processBundle(5, 6, 7), contains(3)); assertThat(tester.processBundle(8, 9), contains(2)); } } }
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 testExtractTotalFlow() throws Exception { DoFnTester<String, KV<String, Integer>> extractFlowInfow = DoFnTester.of(new ExtractFlowInfo()); List<KV<String, Integer>> results = extractFlowInfow.processBundle(INPUT); Assert.assertEquals(1, results.size()); Assert.assertEquals("94", results.get(0).getKey()); Assert.assertEquals(Integer.valueOf(29), results.get(0).getValue()); List<KV<String, Integer>> output = extractFlowInfow.processBundle(""); Assert.assertEquals(0, output.size()); }
/** Example test that tests a specific {@link DoFn}. */ @Test public void testExtractWordsFn() throws Exception { DoFnTester<String, String> extractWordsFn = DoFnTester.of(new ExtractWordsFn()); Assert.assertThat( extractWordsFn.processBundle(" some input words "), CoreMatchers.hasItems("some", "input", "words")); Assert.assertThat(extractWordsFn.processBundle(" "), CoreMatchers.hasItems()); Assert.assertThat( extractWordsFn.processBundle(" some ", " input", " words"), CoreMatchers.hasItems("some", "input", "words")); }
/** Test the {@link ParseEventFn} {@link org.apache.beam.sdk.transforms.DoFn}. */ @Test public void testParseEventFn() throws Exception { DoFnTester<String, GameActionInfo> parseEventFn = DoFnTester.of(new ParseEventFn()); List<GameActionInfo> results = parseEventFn.processBundle(GAME_EVENTS_ARRAY); Assert.assertEquals(8, results.size()); Assert.assertEquals("user0_MagentaKangaroo", results.get(0).getUser()); Assert.assertEquals("MagentaKangaroo", results.get(0).getTeam()); Assert.assertEquals(Integer.valueOf(3), results.get(0).getScore()); }
@Test public void processMultipleBundles() 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")); assertThat(tester.processBundle(5L, 6L, 7L), hasItems("5", "6", "7")); assertThat(tester.processBundle(8L, 9L), hasItems("8", "9")); // peek now returns nothing. assertTrue(tester.peekOutputElements().isEmpty()); } } }
/** A helper method to test {@link SplitQueryFn} to generate the expected number of splits. */ private void testSplitQueryFn( String projectId, String kind, @Nullable String namespace, int expectedNumSplits) throws Exception { Query.Builder query = Query.newBuilder(); query.addKindBuilder().setName(kind); SplitQueryFn splitQueryFn = new SplitQueryFn(V1Options.from(projectId, namespace, null), 0); DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn); List<Query> queries = doFnTester.processBundle(query.build()); assertEquals(expectedNumSplits, queries.size()); }
/** Tests {@link DatastoreV1.Read.SplitQueryFn} when the query has a user specified limit. */ @Test public void testSplitQueryFnWithQueryLimit() throws Exception { Query queryWithLimit = QUERY.toBuilder().setLimit(Int32Value.newBuilder().setValue(1)).build(); SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, 10, mockDatastoreFactory); DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn); doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE); List<Query> queries = doFnTester.processBundle(queryWithLimit); assertEquals(1, queries.size()); verifyNoMoreInteractions(mockDatastore); verifyNoMoreInteractions(mockQuerySplitter); }
@Test public void testCombineVariantsFn() throws Exception { DoFnTester<KV<KV<String, Long>, Iterable<Variant>>, Variant> fn = DoFnTester.of(new JoinNonVariantSegmentsWithVariants.CombineVariantsFn()); Assert.assertThat(fn.processBundle(KV.of(KV.of("chr7", 200000L), (Iterable<Variant>) Arrays.asList(input))), CoreMatchers.hasItems(expectedSnp1, expectedSnp2, expectedInsert)); }
@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 testExtractEventDataFn() throws Exception { DoFnTester<TableRow, KV<String, String>> extractEventDataFn = DoFnTester.of(new ExtractEventDataFn()); List<KV<String, String>> results = extractEventDataFn.processBundle(EVENTS); Assert.assertThat(results, CoreMatchers.hasItem(kv1)); Assert.assertThat(results, CoreMatchers.hasItem(kv2)); }
@Test public void testExtractCountryInfoFn() throws Exception { DoFnTester<TableRow, KV<String, String>> extractCountryInfoFn = DoFnTester.of(new ExtractCountryInfoFn()); List<KV<String, String>> results = extractCountryInfoFn.processBundle(CCS); Assert.assertThat(results, CoreMatchers.hasItem(kv3)); Assert.assertThat(results, CoreMatchers.hasItem(kv4)); }
@Test public void testFormatShakespeareOutputFn() throws Exception { DoFnTester<KV<String, String>, TableRow> formatShakespeareOutputFn = DoFnTester.of(new FormatShakespeareOutputFn()); List<TableRow> results = formatShakespeareOutputFn.processBundle(COMBINED_TUPLES_ARRAY); Assert.assertThat(results, CoreMatchers.hasItem(resultRow1)); Assert.assertThat(results, CoreMatchers.hasItem(resultRow2)); } }
@Test public void testExtractLargeWordsFn() throws Exception { DoFnTester<TableRow, KV<String, String>> extractLargeWordsFn = DoFnTester.of(new ExtractLargeWordsFn()); List<KV<String, String>> results = extractLargeWordsFn.processBundle(ROWS_ARRAY); Assert.assertThat(results, CoreMatchers.hasItem(tuple1)); Assert.assertThat(results, CoreMatchers.hasItem(tuple2)); Assert.assertThat(results, CoreMatchers.hasItem(tuple3)); }
@Test public void testBatchViewAsSingletonToIsmRecord() throws Exception { DoFnTester< KV<Integer, Iterable<KV<GlobalWindow, WindowedValue<String>>>>, IsmRecord<WindowedValue<String>>> doFnTester = DoFnTester.of( new BatchViewOverrides.BatchViewAsSingleton.IsmRecordForSingularValuePerWindowDoFn< String, GlobalWindow>(GlobalWindow.Coder.INSTANCE)); assertThat( doFnTester.processBundle( ImmutableList.of( KV.of( 0, ImmutableList.of(KV.of(GlobalWindow.INSTANCE, valueInGlobalWindow("a")))))), contains(IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE), valueInGlobalWindow("a")))); }
@Test public void testBatchViewAsListToIsmRecordForGlobalWindow() throws Exception { DoFnTester<String, IsmRecord<WindowedValue<String>>> doFnTester = DoFnTester.of( new BatchViewOverrides.BatchViewAsList.ToIsmRecordForGlobalWindowDoFn<String>()); // The order of the output elements is important relative to processing order assertThat( doFnTester.processBundle(ImmutableList.of("a", "b", "c")), contains( IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE, 0L), valueInGlobalWindow("a")), IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE, 1L), valueInGlobalWindow("b")), IsmRecord.of(ImmutableList.of(GlobalWindow.INSTANCE, 2L), valueInGlobalWindow("c")))); }
@Test public void testBinVariantsFn() throws Exception { DoFnTester<Variant, KV<KV<String, Long>, Variant>> binVariantsFn = DoFnTester.of(new JoinNonVariantSegmentsWithVariants.BinShuffleAndCombineTransform.BinVariantsFn()); List<KV<KV<String, Long>, Variant>> binVariantsOutput = binVariantsFn.processBundle(input); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 200000L), snp1))); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 200000L), snp2))); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 200000L), insert))); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 199000L), blockRecord1))); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 200000L), blockRecord1))); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 201000L), blockRecord1))); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 202000L), blockRecord1))); assertThat(binVariantsOutput, CoreMatchers.hasItem(KV.of(KV.of("chr7", 200000L), blockRecord2))); assertEquals(8, binVariantsOutput.size()); }