Examples with PubsubTestClientFactory org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory used on opensource projects

Example 1 with PubsubTestClientFactory

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory in project beam by apache.

the class PubsubUnboundedSinkTest method sendMoreThanOneBatchByByteSize.

@Test
public void sendMoreThanOneBatchByByteSize() throws IOException {
    List<OutgoingMessage> outgoing = new ArrayList<>();
    List<String> data = new ArrayList<>();
    int batchSize = 100;
    int batchBytes = 10;
    int n = 0;
    while (n < batchBytes * 10) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < batchBytes; i++) {
            sb.append(String.valueOf(n));
        }
        String str = sb.toString();
        outgoing.add(new OutgoingMessage(str.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(str)));
        data.add(str);
        n += str.length();
    }
    try (PubsubTestClientFactory factory = PubsubTestClient.createFactoryForPublish(TOPIC, outgoing, ImmutableList.<OutgoingMessage>of())) {
        PubsubUnboundedSink sink = new PubsubUnboundedSink(factory, StaticValueProvider.of(TOPIC), TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, NUM_SHARDS, batchSize, batchBytes, Duration.standardSeconds(2), RecordIdMethod.DETERMINISTIC);
        p.apply(Create.of(data)).apply(ParDo.of(new Stamp())).apply(sink);
        p.run();
    }
// The PubsubTestClientFactory will assert fail on close if the actual published
// message does not match the expected publish message.
}

Example 2 with PubsubTestClientFactory

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory in project beam by apache.

the class PubsubTestClientTest method pullOneMessage.

@Test
public void pullOneMessage() throws IOException {
    final AtomicLong now = new AtomicLong();
    Clock clock = new Clock() {

        @Override
        public long currentTimeMillis() {
            return now.get();
        }
    };
    IncomingMessage expectedIncomingMessage = new IncomingMessage(DATA.getBytes(), null, MESSAGE_TIME, REQ_TIME, ACK_ID, MESSAGE_ID);
    try (PubsubTestClientFactory factory = PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S, Lists.newArrayList(expectedIncomingMessage))) {
        try (PubsubTestClient client = (PubsubTestClient) factory.newClient(null, null, null)) {
            now.set(REQ_TIME);
            client.advance();
            List<IncomingMessage> incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
            assertEquals(1, incomingMessages.size());
            assertEquals(expectedIncomingMessage, incomingMessages.get(0));
            // Timeout on ACK.
            now.addAndGet((ACK_TIMEOUT_S + 10) * 1000);
            client.advance();
            incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
            assertEquals(1, incomingMessages.size());
            assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
            now.addAndGet(10 * 1000);
            client.advance();
            // Extend ack
            client.modifyAckDeadline(SUBSCRIPTION, ImmutableList.of(ACK_ID), 20);
            // Timeout on extended ACK
            now.addAndGet(30 * 1000);
            client.advance();
            incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
            assertEquals(1, incomingMessages.size());
            assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
            // Extend ack
            client.modifyAckDeadline(SUBSCRIPTION, ImmutableList.of(ACK_ID), 20);
            // Ack
            now.addAndGet(15 * 1000);
            client.advance();
            client.acknowledge(SUBSCRIPTION, ImmutableList.of(ACK_ID));
        }
    }
}

Example 3 with PubsubTestClientFactory

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory in project beam by apache.

the class PubsubUnboundedSinkTest method sendOneMessage.

@Test
public void sendOneMessage() throws IOException {
    List<OutgoingMessage> outgoing = ImmutableList.of(new OutgoingMessage(DATA.getBytes(), ATTRIBUTES, TIMESTAMP, getRecordId(DATA)));
    int batchSize = 1;
    int batchBytes = 1;
    try (PubsubTestClientFactory factory = PubsubTestClient.createFactoryForPublish(TOPIC, outgoing, ImmutableList.<OutgoingMessage>of())) {
        PubsubUnboundedSink sink = new PubsubUnboundedSink(factory, StaticValueProvider.of(TOPIC), TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, NUM_SHARDS, batchSize, batchBytes, Duration.standardSeconds(2), RecordIdMethod.DETERMINISTIC);
        p.apply(Create.of(ImmutableList.of(DATA))).apply(ParDo.of(new Stamp(ATTRIBUTES))).apply(sink);
        p.run();
    }
// The PubsubTestClientFactory will assert fail on close if the actual published
// message does not match the expected publish message.
}

Example 4 with PubsubTestClientFactory

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory in project beam by apache.

the class PubsubUnboundedSinkTest method sendMoreThanOneBatchByNumMessages.

@Test
public void sendMoreThanOneBatchByNumMessages() throws IOException {
    List<OutgoingMessage> outgoing = new ArrayList<>();
    List<String> data = new ArrayList<>();
    int batchSize = 2;
    int batchBytes = 1000;
    for (int i = 0; i < batchSize * 10; i++) {
        String str = String.valueOf(i);
        outgoing.add(new OutgoingMessage(str.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(str)));
        data.add(str);
    }
    try (PubsubTestClientFactory factory = PubsubTestClient.createFactoryForPublish(TOPIC, outgoing, ImmutableList.<OutgoingMessage>of())) {
        PubsubUnboundedSink sink = new PubsubUnboundedSink(factory, StaticValueProvider.of(TOPIC), TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, NUM_SHARDS, batchSize, batchBytes, Duration.standardSeconds(2), RecordIdMethod.DETERMINISTIC);
        p.apply(Create.of(data)).apply(ParDo.of(new Stamp())).apply(sink);
        p.run();
    }
// The PubsubTestClientFactory will assert fail on close if the actual published
// message does not match the expected publish message.
}

Example 5 with PubsubTestClientFactory

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory in project beam by apache.

the class PubsubUnboundedSinkTest method sendOneMessageWithoutAttributes.

@Test
public void sendOneMessageWithoutAttributes() throws IOException {
    List<OutgoingMessage> outgoing = ImmutableList.of(new OutgoingMessage(DATA.getBytes(), null, /* attributes */
    TIMESTAMP, getRecordId(DATA)));
    try (PubsubTestClientFactory factory = PubsubTestClient.createFactoryForPublish(TOPIC, outgoing, ImmutableList.<OutgoingMessage>of())) {
        PubsubUnboundedSink sink = new PubsubUnboundedSink(factory, StaticValueProvider.of(TOPIC), TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, NUM_SHARDS, 1, /* batchSize */
        1, /* batchBytes */
        Duration.standardSeconds(2), RecordIdMethod.DETERMINISTIC);
        p.apply(Create.of(ImmutableList.of(DATA))).apply(ParDo.of(new Stamp(null))).apply(sink);
        p.run();
    }
// The PubsubTestClientFactory will assert fail on close if the actual published
// message does not match the expected publish message.
}