use of com.google.cloud.pubsublite.SequencedMessage in project java-pubsublite-spark by googleapis.
the class PslMicroBatchInputPartitionReaderTest method testPartitionReaderNewMessageExceedsRange.
@Test
public void testPartitionReaderNewMessageExceedsRange() throws Exception {
long endOffset = 14L;
createReader(endOffset);
SequencedMessage message1 = newMessage(10L);
SequencedMessage message2 = newMessage(endOffset + 1);
// Multiple get w/o next will return same msg.
when(subscriber.onData()).thenReturn(ApiFutures.immediateFuture(null));
when(subscriber.messageIfAvailable()).thenReturn(Optional.of(message1));
assertThat(reader.next()).isTrue();
verifyInternalRow(reader.get(), 10L);
verifyInternalRow(reader.get(), 10L);
// Next will advance to next msg, and recognize it's out of the batch range.
when(subscriber.onData()).thenReturn(ApiFutures.immediateFuture(null));
when(subscriber.messageIfAvailable()).thenReturn(Optional.of(message2));
assertThat(reader.next()).isFalse();
}
use of com.google.cloud.pubsublite.SequencedMessage in project java-pubsublite-kafka by googleapis.
the class RecordTransformsTest method subscribeTransform.
@Test
public void subscribeTransform() {
SequencedMessage sequencedMessage = SequencedMessage.of(MESSAGE, Timestamp.newBuilder().setNanos(12345).build(), example(Offset.class), 123L);
ConsumerRecord<byte[], byte[]> record = RecordTransforms.fromMessage(sequencedMessage, example(TopicPath.class), example(Partition.class));
assertThat(record.key()).isEqualTo("abc".getBytes());
assertThat(record.value()).isEqualTo("def".getBytes());
assertThat(record.timestampType()).isEqualTo(TimestampType.CREATE_TIME);
assertThat(record.timestamp()).isEqualTo(1001L);
ImmutableListMultimap.Builder<String, ByteString> headers = ImmutableListMultimap.builder();
record.headers().forEach(header -> headers.put(header.key(), ByteString.copyFrom(header.value())));
assertThat(headers.build()).isEqualTo(MESSAGE.attributes());
assertThat(record.offset()).isEqualTo(example(Offset.class).value());
assertThat(record.topic()).isEqualTo(example(TopicPath.class).toString());
assertThat(record.partition()).isEqualTo(example(Partition.class).value());
}
use of com.google.cloud.pubsublite.SequencedMessage in project java-pubsublite by googleapis.
the class MessageTransformsTest method wrappedSubscribeTransformerSetsIdFailure.
@Test
public void wrappedSubscribeTransformerSetsIdFailure() throws ApiException {
MessageTransformer<SequencedMessage, PubsubMessage> mockTransformer = mock(MessageTransformer.class);
MessageTransformer<SequencedMessage, PubsubMessage> wrapped = MessageTransforms.addIdCpsSubscribeTransformer(example(Partition.class), mockTransformer);
when(mockTransformer.transform(any())).thenReturn(PubsubMessage.newBuilder().setMessageId("3").build());
ApiException e = assertThrows(ApiException.class, () -> wrapped.transform(SequencedMessage.of(Message.builder().setAttributes(ImmutableListMultimap.<String, ByteString>builder().put("abc", ByteString.EMPTY).put("def", ByteString.copyFromUtf8("hij")).build()).setData(ByteString.copyFrom(notUtf8Array)).setEventTime(Timestamps.fromNanos(10)).setKey(ByteString.copyFromUtf8("some_key")).build(), Timestamps.fromSeconds(5), Offset.of(7), 2)));
assertThat(e.getStatusCode().getCode()).isEqualTo(Code.INVALID_ARGUMENT);
}
use of com.google.cloud.pubsublite.SequencedMessage in project java-pubsublite by googleapis.
the class SinglePartitionSubscriberTest method multiMessageAck.
@Test
public void multiMessageAck() throws CheckedApiException {
Runnable ack1 = mock(Runnable.class);
Runnable ack2 = mock(Runnable.class);
long bytes2 = 111;
SequencedMessage message2 = SequencedMessage.fromProto(MESSAGE.toProto().toBuilder().setSizeBytes(bytes2).setPublishTime(Timestamps.fromMillis(System.currentTimeMillis())).setCursor(Cursor.newBuilder().setOffset(OFFSET.value() + 1)).build());
when(ackSetTracker.track(MESSAGE)).thenReturn(ack1);
when(ackSetTracker.track(message2)).thenReturn(ack2);
subscriber.onMessages(ImmutableList.of(MESSAGE, message2));
verify(ackSetTracker).track(MESSAGE);
verify(ackSetTracker).track(message2);
verify(receiver).receiveMessage(eq(transformer.transform(MESSAGE)), ackConsumerCaptor.capture());
verify(receiver).receiveMessage(eq(transformer.transform(message2)), ackConsumerCaptor.capture());
ackConsumerCaptor.getAllValues().get(1).ack();
verify(ack2).run();
verify(wireSubscriber).allowFlow(FlowControlRequest.newBuilder().setAllowedMessages(1).setAllowedBytes(bytes2).build());
ackConsumerCaptor.getAllValues().get(0).ack();
verify(ack1).run();
verify(wireSubscriber).allowFlow(FlowControlRequest.newBuilder().setAllowedMessages(1).setAllowedBytes(BYTE_SIZE).build());
}
use of com.google.cloud.pubsublite.SequencedMessage in project java-pubsublite by googleapis.
the class FlowControlBatcherTest method onMessagesDecrementsClientTokens.
@Test
public void onMessagesDecrementsClientTokens() throws CheckedApiException {
FlowControlRequest clientFlowRequest = FlowControlRequest.newBuilder().setAllowedBytes(500).setAllowedMessages(10).build();
batcher.onClientFlowRequest(clientFlowRequest);
ImmutableList<SequencedMessage> messages = ImmutableList.of(SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(0), 100), SequencedMessage.of(Message.builder().build(), Timestamps.EPOCH, Offset.of(1), 150));
batcher.onMessages(messages);
assertThat(batcher.releasePendingRequest().get()).isEqualTo(clientFlowRequest);
FlowControlRequest expectedRequestForRestart = FlowControlRequest.newBuilder().setAllowedBytes(250).setAllowedMessages(8).build();
assertThat(batcher.requestForRestart().get()).isEqualTo(expectedRequestForRestart);
}
Aggregations