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));
}
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();
}
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));
}
Aggregations