Examples with PubsubSource org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubSource used on opensource projects

Example 1 with PubsubSource

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubSource in project beam by apache.

the class PubsubUnboundedSourceTest method setupOneMessage.

private void setupOneMessage(Iterable<IncomingMessage> incoming) {
    now = new AtomicLong(REQ_TIME);
    clock = new Clock() {

        @Override
        public long currentTimeMillis() {
            return now.get();
        }
    };
    factory = PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S, incoming);
    PubsubUnboundedSource source = new PubsubUnboundedSource(clock, factory, null, null, StaticValueProvider.of(SUBSCRIPTION), TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, true);
    primSource = new PubsubSource(source);
}

Example 2 with PubsubSource

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubSource in project beam by apache.

the class PubsubUnboundedSourceTest method noSubscriptionNoSplitGeneratesSubscription.

@Test
public void noSubscriptionNoSplitGeneratesSubscription() throws Exception {
    TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
    factory = PubsubTestClient.createFactoryForCreateSubscription();
    PubsubUnboundedSource source = new PubsubUnboundedSource(factory, StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")), StaticValueProvider.of(topicPath), null, /* subscription */
    null, /* timestampLabel */
    null, /* idLabel */
    false);
    assertThat(source.getSubscription(), nullValue());
    assertThat(source.getSubscription(), nullValue());
    PipelineOptions options = PipelineOptionsFactory.create();
    PubsubSource actualSource = new PubsubSource(source);
    PubsubReader reader = actualSource.createReader(options, null);
    SubscriptionPath createdSubscription = reader.subscription;
    assertThat(createdSubscription, not(nullValue()));
    PubsubCheckpoint checkpoint = reader.getCheckpointMark();
    assertThat(checkpoint.subscriptionPath, equalTo(createdSubscription.getPath()));
    checkpoint.finalizeCheckpoint();
    PubsubCheckpoint deserCheckpoint = CoderUtils.clone(actualSource.getCheckpointMarkCoder(), checkpoint);
    assertThat(checkpoint.subscriptionPath, not(nullValue()));
    assertThat(checkpoint.subscriptionPath, equalTo(deserCheckpoint.subscriptionPath));
    PubsubReader readerFromOriginal = actualSource.createReader(options, checkpoint);
    PubsubReader readerFromDeser = actualSource.createReader(options, deserCheckpoint);
    assertThat(readerFromOriginal.subscription, equalTo(createdSubscription));
    assertThat(readerFromDeser.subscription, equalTo(createdSubscription));
}

Example 3 with PubsubSource

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubSource in project beam by apache.

the class PubsubUnboundedSourceTest method noSubscriptionSplitGeneratesSubscription.

@Test
public void noSubscriptionSplitGeneratesSubscription() throws Exception {
    TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
    factory = PubsubTestClient.createFactoryForCreateSubscription();
    PubsubUnboundedSource source = new PubsubUnboundedSource(factory, StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")), StaticValueProvider.of(topicPath), null, /* subscription */
    null, /* timestampLabel */
    null, /* idLabel */
    false);
    assertThat(source.getSubscription(), nullValue());
    assertThat(source.getSubscription(), nullValue());
    PipelineOptions options = PipelineOptionsFactory.create();
    List<PubsubSource> splits = (new PubsubSource(source)).split(3, options);
    // We have at least one returned split
    assertThat(splits, hasSize(greaterThan(0)));
    for (PubsubSource split : splits) {
        // Each split is equal
        assertThat(split, equalTo(splits.get(0)));
    }
    assertThat(splits.get(0).subscriptionPath, not(nullValue()));
}