@Before public void setup() { p = TestPipeline.create(); visitor = KeyedPValueTrackingVisitor.create(); }
@Test public void testCreationUsingDefaults() { assertNotNull(pipeline); assertNotNull(TestPipeline.create()); }
@Test public void testGetJobMetricsThatFailsForException() throws Exception { DataflowPipelineJob job = spy(new DataflowPipelineJob(mockClient, "test-job", options, null)); Pipeline p = TestPipeline.create(options); p.apply(Create.of(1, 2, 3)); when(mockClient.getJobMetrics(anyString())).thenThrow(new IOException()); TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); assertNull(runner.getJobMetrics(job)); }
@Test public void testBatchOnCreateMatcher() throws Exception { Pipeline p = TestPipeline.create(options); PCollection<Integer> pc = p.apply(Create.of(1, 2, 3)); PAssert.that(pc).containsInAnyOrder(1, 2, 3); final DataflowPipelineJob mockJob = Mockito.mock(DataflowPipelineJob.class); when(mockJob.getState()).thenReturn(State.DONE); when(mockJob.getProjectId()).thenReturn("test-project"); when(mockJob.getJobId()).thenReturn("test-job"); DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class); when(mockRunner.run(any(Pipeline.class))).thenReturn(mockJob); TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); options.as(TestPipelineOptions.class).setOnCreateMatcher(new TestSuccessMatcher(mockJob, 0)); when(mockClient.getJobMetrics(anyString())) .thenReturn(generateMockMetricResponse(true /* success */, true /* tentative */)); runner.run(p, mockRunner); }
@Test public void testBatchOnSuccessMatcherWhenPipelineSucceeds() throws Exception { Pipeline p = TestPipeline.create(options); PCollection<Integer> pc = p.apply(Create.of(1, 2, 3)); PAssert.that(pc).containsInAnyOrder(1, 2, 3); final DataflowPipelineJob mockJob = Mockito.mock(DataflowPipelineJob.class); when(mockJob.getState()).thenReturn(State.DONE); when(mockJob.getProjectId()).thenReturn("test-project"); when(mockJob.getJobId()).thenReturn("test-job"); DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class); when(mockRunner.run(any(Pipeline.class))).thenReturn(mockJob); TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); options.as(TestPipelineOptions.class).setOnSuccessMatcher(new TestSuccessMatcher(mockJob, 1)); when(mockClient.getJobMetrics(anyString())) .thenReturn(generateMockMetricResponse(true /* success */, true /* tentative */)); runner.run(p, mockRunner); }
@Test public void testCreationNotAsTestRule() { thrown.expect(IllegalStateException.class); thrown.expectMessage("@Rule"); TestPipeline.create().run(); }
@Test public void testCheckingForSuccessSkipsNonTentativeMetrics() throws Exception { DataflowPipelineJob job = spy(new DataflowPipelineJob(mockClient, "test-job", options, null)); Pipeline p = TestPipeline.create(options); PCollection<Integer> pc = p.apply(Create.of(1, 2, 3)); PAssert.that(pc).containsInAnyOrder(1, 2, 3); when(mockClient.getJobMetrics(anyString())) .thenReturn( buildJobMetrics(generateMockMetrics(true /* success */, false /* tentative */))); TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); runner.updatePAssertCount(p); doReturn(State.RUNNING).when(job).getState(); assertThat(runner.checkForPAssertSuccess(job), equalTo(Optional.<Boolean>absent())); }
@Test public void testOrderBy_exception() { thrown.expect(UnsupportedOperationException.class); thrown.expectMessage("`ORDER BY` is only supported for GlobalWindows"); String sql = "INSERT INTO SUB_ORDER_RAM(order_id, site_id) SELECT " + " order_id, COUNT(*) " + "FROM ORDER_DETAILS " + "GROUP BY order_id, TUMBLE(order_time, INTERVAL '1' HOUR)" + "ORDER BY order_id asc limit 11"; TestPipeline pipeline = TestPipeline.create(); compilePipeline(sql, pipeline); } }
@Test public void testRuntimeValueProviderTopic() { TestPipeline pipeline = TestPipeline.create(); ValueProvider<String> topic = pipeline.newProvider("projects/project/topics/topic"); Read<String> pubsubRead = PubsubIO.readStrings().fromTopic(topic); pipeline.apply(pubsubRead); assertThat(pubsubRead.getTopicProvider(), not(nullValue())); assertThat(pubsubRead.getTopicProvider().isAccessible(), is(false)); }
@Test public void testExpandWithDuplicates() { Pipeline p = TestPipeline.create(); PCollection<Long> createOne = p.apply("CreateOne", Create.of(1L, 2L, 3L)); PCollectionList<Long> list = PCollectionList.of(createOne).and(createOne).and(createOne); assertThat(list.expand().values(), containsInAnyOrder(createOne, createOne, createOne)); }
/** * Tests that a tentative {@code true} from metrics indicates that every {@link PAssert} has * succeeded. */ @Test public void testCheckingForSuccessWhenPAssertSucceeds() throws Exception { DataflowPipelineJob job = spy(new DataflowPipelineJob(mockClient, "test-job", options, null)); Pipeline p = TestPipeline.create(options); PCollection<Integer> pc = p.apply(Create.of(1, 2, 3)); PAssert.that(pc).containsInAnyOrder(1, 2, 3); when(mockClient.getJobMetrics(anyString())) .thenReturn(buildJobMetrics(generateMockMetrics(true /* success */, true /* tentative */))); TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); doReturn(State.DONE).when(job).getState(); assertThat(runner.checkForPAssertSuccess(job), equalTo(Optional.of(true))); }
@Test public void testGetJobMetricsThatSucceeds() throws Exception { DataflowPipelineJob job = spy(new DataflowPipelineJob(mockClient, "test-job", options, null)); Pipeline p = TestPipeline.create(options); p.apply(Create.of(1, 2, 3)); when(mockClient.getJobMetrics(anyString())) .thenReturn(generateMockMetricResponse(true /* success */, true /* tentative */)); TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient); JobMetrics metrics = runner.getJobMetrics(job); assertEquals(1, metrics.getMetrics().size()); assertEquals( generateMockMetrics(true /* success */, true /* tentative */), metrics.getMetrics()); }
@Test public void testRuntimeValueProviderSubscription() { TestPipeline pipeline = TestPipeline.create(); ValueProvider<String> subscription = pipeline.newProvider("projects/project/subscriptions/subscription"); Read<String> pubsubRead = PubsubIO.readStrings().fromSubscription(subscription); pipeline.apply(pubsubRead); assertThat(pubsubRead.getSubscriptionProvider(), not(nullValue())); assertThat(pubsubRead.getSubscriptionProvider().isAccessible(), is(false)); }
@Test public void getInputSucceeds() { TestPipeline p = TestPipeline.create(); PCollection<KV<String, Integer>> input = p.apply( Create.of(KV.of("foo", 1)) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); PCollection<KV<String, Iterable<Integer>>> grouped = input.apply(GroupByKey.create()); AppliedPTransform<?, ?, ?> producer = DirectGraphs.getProducer(grouped); PTransformReplacement< PCollection<KV<String, Integer>>, PCollection<KV<String, Iterable<Integer>>>> replacement = factory.getReplacementTransform((AppliedPTransform) producer); assertThat(replacement.getInput(), Matchers.<PCollection<?>>equalTo(input)); } }
/** Creates a simple pipeline with a {@link Combine.PerKey}. */ private static TestPipeline createCombinePerKeyPipeline() { TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false); PCollection<KV<String, Integer>> input = pipeline .apply(Create.of(KV.of("key", 1))) .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); input.apply(Combine.perKey(new SumCombineFn())); return pipeline; }
/** Creates a simple pipeline with a {@link Combine.GroupedValues}. */ private static TestPipeline createCombineGroupedValuesPipeline() { TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false); PCollection<KV<String, Integer>> input = pipeline .apply(Create.of(KV.of("key", 1))) .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); input.apply(GroupByKey.create()).apply(Combine.groupedValues(new SumCombineFn())); return pipeline; }
/** * Test to ensure that {@link ValueProvider} values are not accessed at pipeline construction time * when built with {@link DatastoreV1.Read#withLiteralGqlQuery(String)}. */ @Test public void testRuntimeOptionsNotCalledInApplyGqlQuery() { RuntimeTestOptions options = PipelineOptionsFactory.as(RuntimeTestOptions.class); Pipeline pipeline = TestPipeline.create(options); pipeline .apply( DatastoreIO.v1() .read() .withProjectId(options.getDatastoreProject()) .withLiteralGqlQuery(options.getGqlQuery())) .apply(DatastoreIO.v1().write().withProjectId(options.getDatastoreProject())); }
/** Creates a simple pipeline with a {@link Combine.PerKey} with side inputs. */ private static TestPipeline createCombinePerKeyWithSideInputsPipeline() { TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false); PCollection<KV<String, Integer>> input = pipeline .apply(Create.of(KV.of("key", 1))) .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); PCollection<String> sideInput = pipeline.apply(Create.of("side input")); PCollectionView<String> sideInputView = sideInput.apply(View.asSingleton()); input.apply( Combine.<String, Integer, Integer>perKey(new SumCombineFnWithContext()) .withSideInputs(sideInputView)); return pipeline; }
/** * Test to ensure that {@link ValueProvider} values are not accessed at pipeline construction time * when built with {@link DatastoreV1.Read#withQuery(Query)}. */ @Test public void testRuntimeOptionsNotCalledInApplyQuery() { RuntimeTestOptions options = PipelineOptionsFactory.as(RuntimeTestOptions.class); Pipeline pipeline = TestPipeline.create(options); pipeline .apply( DatastoreIO.v1() .read() .withProjectId(options.getDatastoreProject()) .withQuery(QUERY) .withNamespace(options.getNamespace())) .apply(DatastoreIO.v1().write().withProjectId(options.getDatastoreProject())); }
/** Creates a simple pipeline with a {@link Combine.GroupedValues} with side inputs. */ private static TestPipeline createCombineGroupedValuesWithSideInputsPipeline() { TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false); PCollection<KV<String, Integer>> input = pipeline .apply(Create.of(KV.of("key", 1))) .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); PCollection<String> sideInput = pipeline.apply(Create.of("side input")); PCollectionView<String> sideInputView = sideInput.apply(View.asSingleton()); input .apply(GroupByKey.create()) .apply( Combine.<String, Integer, Integer>groupedValues(new SumCombineFnWithContext()) .withSideInputs(sideInputView)); return pipeline; }