FIX: Fix live client test flakiness

Author: nmacholl

tests/test_live_client.py

@@ -358,8 +358,9 @@ async def test_live_start(
     assert message.start_session
 
 
-def test_live_start_twice(
+async def test_live_start_twice(
     live_client: client.Live,
+    mock_live_server: MockLiveServerInterface,
 ) -> None:
     """
     Test that calling start() twice raises a ValueError.
@@ -373,6 +374,10 @@ def test_live_start_twice(
     # Act
     live_client.start()
 
+    _ = await mock_live_server.wait_for_message_of_type(
+        message_type=gateway.SessionStart,
+    )
+
     # Assert
     with pytest.raises(ValueError):
         live_client.start()
@@ -389,8 +394,9 @@ def test_live_start_before_subscribe(
         live_client.start()
 
 
-def test_live_iteration_after_start(
+async def test_live_iteration_after_start(
     live_client: client.Live,
+    mock_live_server: MockLiveServerInterface,
 ) -> None:
     """
     Test that iterating the Live client after it is started raises a
@@ -405,13 +411,18 @@ def test_live_iteration_after_start(
     # Act
     live_client.start()
 
+    _ = await mock_live_server.wait_for_message_of_type(
+        message_type=gateway.SessionStart,
+    )
+
     # Assert
     with pytest.raises(ValueError):
         iter(live_client)
 
 
-def test_live_async_iteration_after_start(
+async def test_live_async_iteration_after_start(
     live_client: client.Live,
+    mock_live_server: MockLiveServerInterface,
 ) -> None:
     """
     Test that async-iterating the Live client after it is started raises a
@@ -426,6 +437,10 @@ def test_live_async_iteration_after_start(
     # Act
     live_client.start()
 
+    _ = await mock_live_server.wait_for_message_of_type(
+        message_type=gateway.SessionStart,
+    )
+
     # Assert
     with pytest.raises(ValueError):
         live_client.__aiter__()
@@ -759,9 +774,10 @@ def test_live_block_for_close(
     assert not live_client.is_connected()
 
 
-def test_live_block_for_close_timeout(
+async def test_live_block_for_close_timeout(
     live_client: client.Live,
     monkeypatch: pytest.MonkeyPatch,
+    mock_live_server: MockLiveServerInterface,
 ) -> None:
     """
     Test that block_for_close terminates the session when the timeout is
@@ -778,15 +794,20 @@ def test_live_block_for_close_timeout(
     )
 
     # Act, Assert
+    _ = await mock_live_server.wait_for_message_of_type(
+        message_type=gateway.SubscriptionRequest,
+    )
+
     live_client.block_for_close(timeout=0)
     live_client.terminate.assert_called_once()  # type: ignore
 
 
 @pytest.mark.usefixtures("mock_live_server")
-def test_live_block_for_close_timeout_stream(
+async def test_live_block_for_close_timeout_stream(
     live_client: client.Live,
     monkeypatch: pytest.MonkeyPatch,
     tmp_path: pathlib.Path,
+    mock_live_server: MockLiveServerInterface,
 ) -> None:
     """
     Test that block_for_close flushes user streams on timeout.
@@ -805,6 +826,10 @@ def test_live_block_for_close_timeout_stream(
     live_client.add_stream(stream)
 
     # Act, Assert
+    _ = await mock_live_server.wait_for_message_of_type(
+        message_type=gateway.SubscriptionRequest,
+    )
+
     live_client.block_for_close(timeout=0)
     stream.flush.assert_called()  # type: ignore [attr-defined]
 
@@ -837,6 +862,7 @@ async def test_live_wait_for_close(
 async def test_live_wait_for_close_timeout(
     live_client: client.Live,
     monkeypatch: pytest.MonkeyPatch,
+    mock_live_server: MockLiveServerInterface,
 ) -> None:
     """
     Test that wait_for_close terminates the session when the timeout is
@@ -853,6 +879,11 @@ async def test_live_wait_for_close_timeout(
         symbols="ALL_SYMBOLS",
         start=None,
     )
+
+    _ = await mock_live_server.wait_for_message_of_type(
+        message_type=gateway.SubscriptionRequest,
+    )
+
     await live_client.wait_for_close(timeout=0)
 
     # Assert
@@ -864,6 +895,7 @@ async def test_live_wait_for_close_timeout_stream(
     live_client: client.Live,
     monkeypatch: pytest.MonkeyPatch,
     tmp_path: pathlib.Path,
+    mock_live_server: MockLiveServerInterface,
 ) -> None:
     """
     Test that wait_for_close flushes user streams on timeout.
@@ -883,6 +915,10 @@ async def test_live_wait_for_close_timeout_stream(
     live_client.add_stream(stream)
 
     # Act
+    _ = await mock_live_server.wait_for_message_of_type(
+        message_type=gateway.SubscriptionRequest,
+    )
+
     await live_client.wait_for_close(timeout=0)
 
     # Assert