review fixes

Author: vgvoleg

tests/topics/test_topic_writer.py

@@ -41,13 +41,10 @@ async def test_random_producer_id(self, driver: ydb.aio.Driver, topic_path, topi
         async with driver.topic_client.writer(topic_path) as writer:
             await writer.write(ydb.TopicWriterMessage(data="123".encode()))
 
-        batch = await topic_reader.receive_batch()
+        msg1 = await topic_reader.receive_message()
+        msg2 = await topic_reader.receive_message()
 
-        if len(batch.messages) == 1:
-            batch2 = await topic_reader.receive_batch()
-            batch.messages.extend(batch2.messages)
-
-        assert batch.messages[0].producer_id != batch.messages[1].producer_id
+        assert msg1.producer_id != msg2.producer_id
 
     async def test_auto_flush_on_close(self, driver: ydb.aio.Driver, topic_path):
         async with driver.topic_client.writer(
@@ -80,18 +77,16 @@ async def test_write_multi_message_with_ack(
             assert res1.offset == 0
             assert res2.offset == 1
 
-        batch = await topic_reader.receive_batch()
+        msg1 = await topic_reader.receive_message()
+        msg2 = await topic_reader.receive_message()
 
-        assert batch.messages[0].offset == 0
-        assert batch.messages[0].seqno == 1
-        assert batch.messages[0].data == "123".encode()
+        assert msg1.offset == 0
+        assert msg1.seqno == 1
+        assert msg1.data == "123".encode()
 
-        # # remove second recieve batch when implement batching
-        # # https://github.com/ydb-platform/ydb-python-sdk/issues/142
-        # batch = await topic_reader.receive_batch()
-        assert batch.messages[1].offset == 1
-        assert batch.messages[1].seqno == 2
-        assert batch.messages[1].data == "456".encode()
+        assert msg2.offset == 1
+        assert msg2.seqno == 2
+        assert msg2.data == "456".encode()
 
     @pytest.mark.parametrize(
         "codec",
@@ -189,9 +184,10 @@ def test_random_producer_id(
         with driver_sync.topic_client.writer(topic_path) as writer:
             writer.write(ydb.TopicWriterMessage(data="123".encode()))
 
-        batch = topic_reader_sync.receive_batch()
+        msg1 = topic_reader_sync.receive_message()
+        msg2 = topic_reader_sync.receive_message()
 
-        assert batch.messages[0].producer_id != batch.messages[1].producer_id
+        assert msg1.producer_id != msg2.producer_id
 
     def test_write_multi_message_with_ack(
         self, driver_sync: ydb.Driver, topic_path, topic_reader_sync: ydb.TopicReader
@@ -204,18 +200,16 @@ def test_write_multi_message_with_ack(
                 ]
             )
 
-        batch = topic_reader_sync.receive_batch()
-        if len(batch.messages) == 1:
-            batch2 = topic_reader_sync.receive_batch()
-            batch.messages.extend(batch2.messages)
+        msg1 = topic_reader_sync.receive_message()
+        msg2 = topic_reader_sync.receive_message()
 
-        assert batch.messages[0].offset == 0
-        assert batch.messages[0].seqno == 1
-        assert batch.messages[0].data == "123".encode()
+        assert msg1.offset == 0
+        assert msg1.seqno == 1
+        assert msg1.data == "123".encode()
 
-        assert batch.messages[1].offset == 1
-        assert batch.messages[1].seqno == 2
-        assert batch.messages[1].data == "456".encode()
+        assert msg2.offset == 1
+        assert msg2.seqno == 2
+        assert msg2.data == "456".encode()
 
     @pytest.mark.parametrize(
         "codec",

ydb/_topic_reader/topic_reader_asyncio.py

@@ -264,7 +264,7 @@ class ReaderStream:
 
     _state_changed: asyncio.Event
     _closed: bool
-    _message_batches: typing.Dict[int, datatypes.PublicBatch]
+    _message_batches: typing.Dict[int, datatypes.PublicBatch]  # keys are partition session ID
     _first_error: asyncio.Future[YdbError]
 
     _update_token_interval: Union[int, float]
@@ -360,8 +360,8 @@ async def wait_messages(self):
             self._state_changed.clear()
 
     def _get_first_batch(self) -> typing.Tuple[int, datatypes.PublicBatch]:
-        first_id, batch = self._message_batches.popitem(last=False)
-        return first_id, batch
+        partition_session_id, batch = self._message_batches.popitem(last=False)
+        return partition_session_id, batch
 
     def receive_batch_nowait(self):
         if self._get_first_error():

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -232,7 +232,7 @@ def batch_size():
             return len(stream_reader._message_batches[partition_session_id].messages)
 
         initial_batches = batch_count()
-        initial_batch_size = batch_size() if not new_batch else 0
+        initial_batch_size = 0 if new_batch else batch_size()
 
         stream = stream_reader._stream  # type: StreamMock
         stream.from_server.put_nowait(