Search in sources :

Example 1 with PubSubWritePayload

use of org.apache.beam.model.pipeline.v1.RunnerApi.PubSubWritePayload in project beam by apache.

the class PubSubWritePayloadTranslationTest method testTranslateSinkWithTopicOverridden.

@Test
public void testTranslateSinkWithTopicOverridden() throws Exception {
    ValueProvider<TopicPath> runtimeProvider = pipeline.newProvider(TOPIC);
    PubsubUnboundedSink pubsubUnboundedSinkSink = new PubsubUnboundedSink(null, runtimeProvider, TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, 0, 0, 0, Duration.ZERO, null);
    PubsubSink pubsubSink = new PubsubSink(pubsubUnboundedSinkSink);
    PCollection<byte[]> input = pipeline.apply(Create.of(new byte[0]));
    PDone output = input.apply(pubsubSink);
    AppliedPTransform<?, ?, PubsubSink> appliedPTransform = AppliedPTransform.of("sink", PValues.expandInput(input), PValues.expandOutput(output), pubsubSink, ResourceHints.create(), pipeline);
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.FunctionSpec spec = sinkTranslator.translate(appliedPTransform, components);
    assertEquals(PTransformTranslation.PUBSUB_WRITE, spec.getUrn());
    PubSubWritePayload payload = PubSubWritePayload.parseFrom(spec.getPayload());
    assertEquals(((NestedValueProvider) runtimeProvider).propertyName(), payload.getTopicRuntimeOverridden());
    assertTrue(payload.getTopic().isEmpty());
    assertEquals(TIMESTAMP_ATTRIBUTE, payload.getTimestampAttribute());
    assertEquals(ID_ATTRIBUTE, payload.getIdAttribute());
}
Also used : RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) TopicPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath) PDone(org.apache.beam.sdk.values.PDone) PubsubSink(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSink.PubsubSink) SdkComponents(org.apache.beam.runners.core.construction.SdkComponents) PubSubWritePayload(org.apache.beam.model.pipeline.v1.RunnerApi.PubSubWritePayload) Test(org.junit.Test)

Example 2 with PubSubWritePayload

use of org.apache.beam.model.pipeline.v1.RunnerApi.PubSubWritePayload in project beam by apache.

the class PubSubWritePayloadTranslationTest method testTranslateSinkWithTopic.

@Test
public void testTranslateSinkWithTopic() throws Exception {
    PubsubUnboundedSink pubsubUnboundedSink = new PubsubUnboundedSink(null, StaticValueProvider.of(TOPIC), TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, 0, 0, 0, Duration.ZERO, null);
    PubsubUnboundedSink.PubsubSink pubsubSink = new PubsubSink(pubsubUnboundedSink);
    PCollection<byte[]> input = pipeline.apply(Create.of(new byte[0]));
    PDone output = input.apply(pubsubSink);
    AppliedPTransform<?, ?, PubsubSink> appliedPTransform = AppliedPTransform.of("sink", PValues.expandInput(input), PValues.expandOutput(output), pubsubSink, ResourceHints.create(), pipeline);
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.FunctionSpec spec = sinkTranslator.translate(appliedPTransform, components);
    assertEquals(PTransformTranslation.PUBSUB_WRITE, spec.getUrn());
    PubSubWritePayload payload = PubSubWritePayload.parseFrom(spec.getPayload());
    assertEquals(TOPIC.getFullPath(), payload.getTopic());
    assertTrue(payload.getTopicRuntimeOverridden().isEmpty());
    assertEquals(TIMESTAMP_ATTRIBUTE, payload.getTimestampAttribute());
    assertEquals(ID_ATTRIBUTE, payload.getIdAttribute());
}
Also used : RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) PDone(org.apache.beam.sdk.values.PDone) PubsubSink(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSink.PubsubSink) PubsubSink(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSink.PubsubSink) SdkComponents(org.apache.beam.runners.core.construction.SdkComponents) PubSubWritePayload(org.apache.beam.model.pipeline.v1.RunnerApi.PubSubWritePayload) Test(org.junit.Test)

Aggregations

RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)2 PubSubWritePayload (org.apache.beam.model.pipeline.v1.RunnerApi.PubSubWritePayload)2 SdkComponents (org.apache.beam.runners.core.construction.SdkComponents)2 PubsubSink (org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSink.PubsubSink)2 PDone (org.apache.beam.sdk.values.PDone)2 Test (org.junit.Test)2 TopicPath (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath)1