Examples with IncomingMessage org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage used on opensource projects

Example 6 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage in project beam by apache.

the class PubsubJsonClientTest method pullOneMessageEmptyAttributes.

@Test
public void pullOneMessageEmptyAttributes() throws IOException {
    client = new PubsubJsonClient(null, null, mockPubsub);
    String expectedSubscription = SUBSCRIPTION.getPath();
    PullRequest expectedRequest = new PullRequest().setReturnImmediately(true).setMaxMessages(10);
    PubsubMessage expectedPubsubMessage = new PubsubMessage().setMessageId(MESSAGE_ID).encodeData(DATA.getBytes(StandardCharsets.UTF_8)).setPublishTime(String.valueOf(PUB_TIME));
    ReceivedMessage expectedReceivedMessage = new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
    PullResponse expectedResponse = new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
    when((Object) mockPubsub.projects().subscriptions().pull(expectedSubscription, expectedRequest).execute()).thenReturn(expectedResponse);
    List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
    assertEquals(1, acutalMessages.size());
    IncomingMessage actualMessage = acutalMessages.get(0);
    assertEquals(ACK_ID, actualMessage.ackId());
    assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
    assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch());
    assertEquals(PUB_TIME, actualMessage.timestampMsSinceEpoch());
}

Example 7 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage in project beam by apache.

the class PubsubIOTest method testFailedParseWithDeadLetterConfigured.

@Test
public void testFailedParseWithDeadLetterConfigured() {
    ByteString data = ByteString.copyFrom("Hello, World!".getBytes(StandardCharsets.UTF_8));
    RuntimeException exception = new RuntimeException("Some error message");
    ImmutableList<IncomingMessage> expectedReads = ImmutableList.of(IncomingMessage.of(com.google.pubsub.v1.PubsubMessage.newBuilder().setData(data).build(), 1234L, 0, UUID.randomUUID().toString(), UUID.randomUUID().toString()));
    ImmutableList<OutgoingMessage> expectedWrites = ImmutableList.of(OutgoingMessage.of(com.google.pubsub.v1.PubsubMessage.newBuilder().setData(data).putAttributes("exceptionClassName", exception.getClass().getName()).putAttributes("exceptionMessage", exception.getMessage()).putAttributes("pubsubMessageId", "<null>").build(), 1234L, null));
    clientFactory = PubsubTestClient.createFactoryForPullAndPublish(SUBSCRIPTION, TOPIC, CLOCK, 60, expectedReads, expectedWrites, ImmutableList.of());
    PCollection<String> read = readPipeline.apply(PubsubIO.readStrings().fromSubscription(SUBSCRIPTION.getPath()).withDeadLetterTopic(TOPIC.getPath()).withClock(CLOCK).withClientFactory(clientFactory).withCoderAndParseFn(StringUtf8Coder.of(), SimpleFunction.fromSerializableFunctionWithOutputType(message -> {
        throw exception;
    }, TypeDescriptors.strings())));
    PAssert.that(read).empty();
    readPipeline.run();
}

Example 8 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage in project beam by apache.

the class PubsubTestClientTest method pullOneMessage.

@Test
public void pullOneMessage() throws IOException {
    final AtomicLong now = new AtomicLong();
    Clock clock = now::get;
    IncomingMessage expectedIncomingMessage = IncomingMessage.of(PubsubMessage.newBuilder().setData(ByteString.copyFromUtf8(DATA)).build(), 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(IncomingMessage.of(expectedIncomingMessage.message(), expectedIncomingMessage.timestampMsSinceEpoch(), now.get(), expectedIncomingMessage.ackId(), expectedIncomingMessage.recordId()), 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(IncomingMessage.of(expectedIncomingMessage.message(), expectedIncomingMessage.timestampMsSinceEpoch(), now.get(), expectedIncomingMessage.ackId(), expectedIncomingMessage.recordId()), 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 9 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage in project beam by apache.

the class PubsubTestClientTest method testPullThenPublish.

@Test
public void testPullThenPublish() throws IOException {
    AtomicLong now = new AtomicLong();
    Clock clock = now::get;
    PubsubMessage message = PubsubMessage.newBuilder().setData(ByteString.copyFromUtf8(DATA)).build();
    IncomingMessage expectedIncomingMessage = IncomingMessage.of(message, MESSAGE_TIME, REQ_TIME, ACK_ID, MESSAGE_ID);
    OutgoingMessage expectedOutgoingMessage = OutgoingMessage.of(message, MESSAGE_TIME, MESSAGE_ID);
    try (PubsubTestClientFactory factory = PubsubTestClient.createFactoryForPullAndPublish(SUBSCRIPTION, TOPIC, clock, ACK_TIMEOUT_S, ImmutableList.of(expectedIncomingMessage), ImmutableList.of(expectedOutgoingMessage), ImmutableList.of())) {
        try (PubsubTestClient client = (PubsubTestClient) factory.newClient(null, null, null)) {
            // Pull
            now.set(REQ_TIME);
            client.advance();
            List<IncomingMessage> actualIncomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
            now.addAndGet(ACK_TIMEOUT_S - 10);
            client.advance();
            client.acknowledge(SUBSCRIPTION, ImmutableList.of(ACK_ID));
            assertEquals(1, actualIncomingMessages.size());
            assertEquals(expectedIncomingMessage, actualIncomingMessages.get(0));
            // Publish
            IncomingMessage incomingMessage = actualIncomingMessages.get(0);
            OutgoingMessage actualOutgoingMessage = OutgoingMessage.of(incomingMessage.message(), incomingMessage.timestampMsSinceEpoch(), incomingMessage.recordId());
            client.publish(TOPIC, ImmutableList.of(actualOutgoingMessage));
        }
    }
}

Example 10 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage in project beam by apache.

the class PubsubJsonClientTest method pullOneMessageWithNoData.

@Test
public void pullOneMessageWithNoData() throws IOException {
    String expectedSubscription = SUBSCRIPTION.getPath();
    PullRequest expectedRequest = new PullRequest().setReturnImmediately(true).setMaxMessages(10);
    PubsubMessage expectedPubsubMessage = new PubsubMessage().setMessageId(MESSAGE_ID).setPublishTime(String.valueOf(PUB_TIME)).setAttributes(ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME), ID_ATTRIBUTE, RECORD_ID));
    ReceivedMessage expectedReceivedMessage = new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
    PullResponse expectedResponse = new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
    Mockito.when((Object) mockPubsub.projects().subscriptions().pull(expectedSubscription, expectedRequest).execute()).thenReturn(expectedResponse);
    List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
    assertEquals(1, acutalMessages.size());
    IncomingMessage actualMessage = acutalMessages.get(0);
    assertArrayEquals(new byte[0], actualMessage.message().getData().toByteArray());
}