better tests fot topic_writer/topic_reader

Author: Valeriya Popova

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -125,17 +125,8 @@ def second_partition_session(
 
         return stream_reader_started._partition_sessions[partition_session.id]
 
-    @pytest.fixture()
-    async def stream_reader_started(
-        self, stream, default_reader_settings, request
-    ) -> ReaderStream:
-
-        settings, token_getter = getattr(
-            request, "param", (default_reader_settings, None)
-        )
-        reader = ReaderStream(
-            self.default_reader_reconnector_id, settings, token_getter
-        )
+    async def get_started_reader(self, stream, *args, **kwargs) -> ReaderStream:
+        reader = ReaderStream(self.default_reader_reconnector_id, *args, **kwargs)
         init_message = object()
 
         # noinspection PyTypeChecker
@@ -164,6 +155,12 @@ async def stream_reader_started(
 
         return reader
 
+    @pytest.fixture()
+    async def stream_reader_started(
+        self, stream, default_reader_settings
+    ) -> ReaderStream:
+        return await self.get_started_reader(stream, default_reader_settings)
+
     @pytest.fixture()
     async def stream_reader(self, stream_reader_started: ReaderStream):
         yield stream_reader_started
@@ -1012,21 +1009,16 @@ async def test_receive_batch_nowait(self, stream, stream_reader, partition_sessi
         with pytest.raises(asyncio.QueueEmpty):
             stream.from_client.get_nowait()
 
-    @pytest.mark.parametrize(
-        "stream_reader_started",
-        [
-            (
-                PublicReaderSettings(
-                    consumer="test-consumer",
-                    topic="test-topic",
-                    update_token_interval=0.1,
-                ),
-                lambda: "foo-bar",
-            )
-        ],
-        indirect=True,
-    )
-    async def test_update_token(self, stream, stream_reader_started: ReaderStream):
+    async def test_update_token(self, stream):
+        settings = PublicReaderSettings(
+            consumer="test-consumer",
+            topic="test-topic",
+            update_token_interval=0.1,
+        )
+        reader = await self.get_started_reader(
+            stream, settings, get_token_function=lambda: "foo-bar"
+        )
+
         assert stream.from_client.empty()
 
         expected = StreamReadMessage.FromClient(UpdateTokenRequest(token="foo-bar"))
@@ -1046,7 +1038,7 @@ async def test_update_token(self, stream, stream_reader_started: ReaderStream):
         got = await wait_for_fast(stream.from_client.get())
         assert expected == got
 
-        await stream_reader_started.close()
+        await reader.close()
 
 
 @pytest.mark.asyncio

ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -63,8 +63,8 @@ def stream(self):
         yield stream
         stream.close()
 
-    @pytest.fixture
-    async def writer_and_stream(self, stream, request) -> WriterWithMockedStream:
+    @staticmethod
+    async def get_started_writer(stream, *args, **kwargs) -> WriterAsyncIOStream:
         stream.from_server.put_nowait(
             StreamWriteMessage.InitResponse(
                 last_seq_no=4,
@@ -75,9 +75,7 @@ async def writer_and_stream(self, stream, request) -> WriterWithMockedStream:
             )
         )
 
-        params = getattr(request, "param", ())
-        writer = WriterAsyncIOStream(*params)
-
+        writer = WriterAsyncIOStream(*args, **kwargs)
         await writer._start(
             stream,
             init_message=StreamWriteMessage.InitRequest(
@@ -91,6 +89,11 @@ async def writer_and_stream(self, stream, request) -> WriterWithMockedStream:
             ),
         )
         await stream.from_client.get()
+        return writer
+
+    @pytest.fixture
+    async def writer_and_stream(self, stream) -> WriterWithMockedStream:
+        writer = await self.get_started_writer(stream)
 
         yield TestWriterAsyncIOStream.WriterWithMockedStream(
             stream=stream,
@@ -164,25 +167,23 @@ async def test_write_a_message(self, writer_and_stream: WriterWithMockedStream):
         sent_message = await writer_and_stream.stream.from_client.get()
         assert expected_message == sent_message
 
-    @pytest.mark.parametrize(
-        "writer_and_stream", [(0.1, lambda: "foo-bar")], indirect=True
-    )
-    async def test_update_token(self, writer_and_stream: WriterWithMockedStream):
-        assert writer_and_stream.stream.from_client.empty()
+    async def test_update_token(self, stream: StreamMock):
+        writer = await self.get_started_writer(
+            stream, update_token_interval=0.1, get_token_function=lambda: "foo-bar"
+        )
+        assert stream.from_client.empty()
 
         expected = StreamWriteMessage.FromClient(UpdateTokenRequest(token="foo-bar"))
-        got = await wait_for_fast(writer_and_stream.stream.from_client.get())
+        got = await wait_for_fast(stream.from_client.get())
         assert expected == got, "send update token request"
 
         await asyncio.sleep(0.2)
-        assert (
-            writer_and_stream.stream.from_client.empty()
-        ), "no answer - no new update request"
+        assert stream.from_client.empty(), "no answer - no new update request"
 
-        await writer_and_stream.stream.from_server.put(UpdateTokenResponse())
-        receive_task = asyncio.create_task(writer_and_stream.writer.receive())
+        await stream.from_server.put(UpdateTokenResponse())
+        receive_task = asyncio.create_task(writer.receive())
 
-        got = await wait_for_fast(writer_and_stream.stream.from_client.get())
+        got = await wait_for_fast(stream.from_client.get())
         assert expected == got
 
         receive_task.cancel()