diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java index 0a055a0e3..3fd8099ff 100644 --- a/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java +++ b/topic/src/main/java/tech/ydb/topic/read/impl/ReadPartitionSession.java @@ -92,7 +92,7 @@ public CompletableFuture addBatches(List delivered = Collections.synchronizedList(new ArrayList<>()); + + ReadPartitionSession reader = new ReadPartitionSession("test", session, partition, 0) { + @Override + CompletableFuture handleDataReceivedEvent(DataReceivedEvent event) { + event.getMessages().forEach(message -> delivered.add(message.getOffset())); + return CompletableFuture.completedFuture(null); + } + }; + + CompletableFuture batchesRead = reader.addBatches( + Arrays.asList(emptyProtoBatch(), rawProtoBatch(1), rawProtoBatch(2)) + ); + + batchesRead.get(); + Assert.assertEquals(Arrays.asList(1L, 2L), delivered); + } + + private static YdbTopic.StreamReadMessage.ReadResponse.Batch emptyProtoBatch() { + return YdbTopic.StreamReadMessage.ReadResponse.Batch.newBuilder() + .setCodec(Codec.RAW) + .setProducerId("producer") + .build(); + } + + private static YdbTopic.StreamReadMessage.ReadResponse.Batch rawProtoBatch(long offset) { + return YdbTopic.StreamReadMessage.ReadResponse.Batch.newBuilder() + .setCodec(Codec.RAW) + .setProducerId("producer") + .addMessageData( + YdbTopic.StreamReadMessage.ReadResponse.MessageData.newBuilder() + .setOffset(offset) + .setSeqNo(offset) + .setData(ByteString.copyFromUtf8("data")) + .build() + ) + .build(); + } +}