@Test public void registerTransformNoChildren() throws IOException { Create.Values<Integer> create = Create.of(1, 2, 3); PCollection<Integer> pt = pipeline.apply(create); String userName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); String componentName = components.registerPTransform(transform, Collections.emptyList()); assertThat(componentName, equalTo(userName)); assertThat(components.getExistingPTransformId(transform), equalTo(componentName)); }
private static AppliedPTransform<?, ?, ?> generateSequence(Pipeline pipeline) { GenerateSequence sequence = GenerateSequence.from(0); PCollection<Long> pcollection = pipeline.apply(sequence); return AppliedPTransform.of( "Count", pipeline.begin().expand(), pcollection.expand(), sequence, pipeline); }
private static AppliedPTransform<?, ?, ?> rawPTransformWithNullSpec(Pipeline pipeline) { PTransformTranslation.RawPTransform<PBegin, PDone> rawPTransform = new PTransformTranslation.RawPTransform<PBegin, PDone>() { @Override public String getUrn() { return "fake/urn"; } @Nullable @Override public RunnerApi.FunctionSpec getSpec() { return null; } }; return AppliedPTransform.<PBegin, PDone, PTransform<PBegin, PDone>>of( "RawPTransformWithNoSpec", pipeline.begin().expand(), PDone.in(pipeline).expand(), rawPTransform, pipeline); }
private static AppliedPTransform<?, ?, ?> read(Pipeline pipeline) { Read.Unbounded<Long> transform = Read.from(CountingSource.unbounded()); PCollection<Long> pcollection = pipeline.apply(transform); return AppliedPTransform.of( "ReadTheCount", pipeline.begin().expand(), pcollection.expand(), transform, pipeline); }
/** Tests that trying to register a transform which has unregistered children throws. */ @Test public void registerTransformWithUnregisteredChildren() throws IOException { Create.Values<Long> create = Create.of(1L, 2L, 3L); GenerateSequence createChild = GenerateSequence.from(0); PCollection<Long> pt = pipeline.apply(create); String userName = "my_transform"; String childUserName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); AppliedPTransform<?, ?, ?> childTransform = AppliedPTransform.of( childUserName, pipeline.begin().expand(), pt.expand(), createChild, pipeline); thrown.expect(IllegalArgumentException.class); thrown.expectMessage(childTransform.toString()); components.registerPTransform(transform, Collections.singletonList(childTransform)); }
@Test public void registerTransformEmptyFullName() throws IOException { Create.Values<Integer> create = Create.of(1, 2, 3); PCollection<Integer> pt = pipeline.apply(create); AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of("", pipeline.begin().expand(), pt.expand(), create, pipeline); thrown.expect(IllegalArgumentException.class); thrown.expectMessage(transform.toString()); components.getExistingPTransformId(transform); }
@Test public void registerTransformAfterChildren() throws IOException { Create.Values<Long> create = Create.of(1L, 2L, 3L); GenerateSequence createChild = GenerateSequence.from(0); PCollection<Long> pt = pipeline.apply(create); String userName = "my_transform"; String childUserName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); AppliedPTransform<?, ?, ?> childTransform = AppliedPTransform.of( childUserName, pipeline.begin().expand(), pt.expand(), createChild, pipeline); String childId = components.registerPTransform(childTransform, Collections.emptyList()); String parentId = components.registerPTransform(transform, Collections.singletonList(childTransform)); Components components = this.components.toComponents(); assertThat(components.getTransformsOrThrow(parentId).getSubtransforms(0), equalTo(childId)); assertThat(components.getTransformsOrThrow(childId).getSubtransformsCount(), equalTo(0)); }
@Test public void registerTransformNullComponents() throws IOException { Create.Values<Integer> create = Create.of(1, 2, 3); PCollection<Integer> pt = pipeline.apply(create); String userName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); thrown.expect(NullPointerException.class); thrown.expectMessage("child nodes may not be null"); components.registerPTransform(transform, null); }
@Test public void testRegistrarEncodedProto() throws Exception { PCollection<String> output = p.apply(testStream); AppliedPTransform<PBegin, PCollection<String>, TestStream<String>> appliedTestStream = AppliedPTransform.of("fakeName", PBegin.in(p).expand(), output.expand(), testStream, p); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.FunctionSpec spec = PTransformTranslation.toProto(appliedTestStream, components).getSpec(); assertThat(spec.getUrn(), equalTo(TEST_STREAM_TRANSFORM_URN)); RunnerApi.TestStreamPayload payload = TestStreamPayload.parseFrom(spec.getPayload()); verifyTestStreamEncoding( testStream, payload, RehydratedComponents.forComponents(components.toComponents())); }