Search in sources :

Example 1 with Message

use of com.google.cloud.pubsublite.Message in project beam by apache.

the class PubsubLiteSinkTest method manyMessagePublishes.

@Test
public void manyMessagePublishes() throws Exception {
    Message message1 = Message.builder().build();
    Message message2 = Message.builder().setKey(ByteString.copyFromUtf8("abc")).build();
    when(publisher.publish(message1)).thenReturn(ApiFutures.immediateFuture(MessageMetadata.of(Partition.of(1), Offset.of(2))));
    when(publisher.publish(message2)).thenReturn(ApiFutures.immediateFuture(MessageMetadata.of(Partition.of(85), Offset.of(3))));
    runWith(message1, message2);
    verify(publisher, times(2)).publish(publishedMessageCaptor.capture());
    assertThat(publishedMessageCaptor.getAllValues(), containsInAnyOrder(message1, message2));
}
Also used : Message(com.google.cloud.pubsublite.Message) Test(org.junit.Test)

Example 2 with Message

use of com.google.cloud.pubsublite.Message in project beam by apache.

the class PubsubLiteSinkTest method exceptionMixedWithOK.

@Test
public void exceptionMixedWithOK() throws Exception {
    Message message1 = Message.builder().build();
    Message message2 = Message.builder().setKey(ByteString.copyFromUtf8("abc")).build();
    Message message3 = Message.builder().setKey(ByteString.copyFromUtf8("def")).build();
    SettableApiFuture<MessageMetadata> future1 = SettableApiFuture.create();
    SettableApiFuture<MessageMetadata> future2 = SettableApiFuture.create();
    SettableApiFuture<MessageMetadata> future3 = SettableApiFuture.create();
    CountDownLatch startedLatch = new CountDownLatch(3);
    when(publisher.publish(message1)).then(invocation -> {
        startedLatch.countDown();
        return future1;
    });
    when(publisher.publish(message2)).then(invocation -> {
        startedLatch.countDown();
        return future2;
    });
    when(publisher.publish(message3)).then(invocation -> {
        startedLatch.countDown();
        return future3;
    });
    ExecutorService exec = Executors.newCachedThreadPool();
    exec.execute(() -> {
        try {
            startedLatch.await();
            future1.set(MessageMetadata.of(Partition.of(1), Offset.of(2)));
            future2.setException(new CheckedApiException(Code.INTERNAL).underlying);
            future3.set(MessageMetadata.of(Partition.of(1), Offset.of(3)));
        } catch (InterruptedException e) {
            fail();
            throw new RuntimeException(e);
        }
    });
    PipelineExecutionException e = assertThrows(PipelineExecutionException.class, () -> runWith(message1, message2, message3));
    verify(publisher, times(3)).publish(publishedMessageCaptor.capture());
    assertThat(publishedMessageCaptor.getAllValues(), containsInAnyOrder(message1, message2, message3));
    Optional<CheckedApiException> statusOr = ExtractStatus.extract(e.getCause());
    assertTrue(statusOr.isPresent());
    assertThat(statusOr.get().code(), equalTo(Code.INTERNAL));
    exec.shutdownNow();
}
Also used : MessageMetadata(com.google.cloud.pubsublite.MessageMetadata) Message(com.google.cloud.pubsublite.Message) PipelineExecutionException(org.apache.beam.sdk.Pipeline.PipelineExecutionException) ExecutorService(java.util.concurrent.ExecutorService) CheckedApiException(com.google.cloud.pubsublite.internal.CheckedApiException) CountDownLatch(java.util.concurrent.CountDownLatch) Test(org.junit.Test)

Example 3 with Message

use of com.google.cloud.pubsublite.Message in project beam by apache.

the class PubsubLiteSinkTest method singleExceptionWhenProcessing.

@Test
public void singleExceptionWhenProcessing() {
    Message message1 = Message.builder().build();
    when(publisher.publish(message1)).thenReturn(ApiFutures.immediateFailedFuture(new CheckedApiException(Code.INTERNAL).underlying));
    PipelineExecutionException e = assertThrows(PipelineExecutionException.class, () -> runWith(message1));
    verify(publisher).publish(message1);
    Optional<CheckedApiException> statusOr = ExtractStatus.extract(e.getCause());
    assertTrue(statusOr.isPresent());
    assertThat(statusOr.get().code(), equalTo(Code.INTERNAL));
}
Also used : Message(com.google.cloud.pubsublite.Message) PipelineExecutionException(org.apache.beam.sdk.Pipeline.PipelineExecutionException) CheckedApiException(com.google.cloud.pubsublite.internal.CheckedApiException) Test(org.junit.Test)

Aggregations

Message (com.google.cloud.pubsublite.Message)3 Test (org.junit.Test)3 CheckedApiException (com.google.cloud.pubsublite.internal.CheckedApiException)2 PipelineExecutionException (org.apache.beam.sdk.Pipeline.PipelineExecutionException)2 MessageMetadata (com.google.cloud.pubsublite.MessageMetadata)1 CountDownLatch (java.util.concurrent.CountDownLatch)1 ExecutorService (java.util.concurrent.ExecutorService)1