FIX: Live client test flakiness

Author: nmacholl

databento/live/dbn.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 import queue
+import threading
 from concurrent import futures
 from typing import Callable, NewType, Optional, Union
 
@@ -57,6 +58,7 @@ def __init__(
         self._dbn_queue: Optional[DBNQueue] = None
         self._decoder: databento_dbn.DbnDecoder = databento_dbn.DbnDecoder()
         self._disconnected: "asyncio.Future[None]" = loop.create_future()
+        self._transport_lock = threading.Lock()
         self._transport = transport
 
     def __aiter__(self) -> "DBNProtocol":
@@ -90,7 +92,8 @@ async def __anext__(self) -> DBNStruct:
                         "resuming reading with %d pending records",
                         self._dbn_queue.qsize(),
                     )
-                    self._transport.resume_reading()
+                    with self._transport_lock:
+                        self._transport.resume_reading()
 
         raise StopAsyncIteration()
 
@@ -118,7 +121,8 @@ def __next__(self) -> DBNStruct:
                         "resuming reading with %d pending records",
                         self._dbn_queue.qsize(),
                     )
-                    self._transport.resume_reading()
+                    with self._transport_lock:
+                        self._transport.resume_reading()
 
         raise StopIteration()
 
@@ -171,7 +175,8 @@ def buffer_updated(self, nbytes: int) -> None:
             self._decoder.write(record_bytes)
         except ValueError:
             logger.critical("could not write to dbn decoder")
-            self._transport.close()
+            with self._transport_lock:
+                self._transport.close()
 
         try:
             records = self._decoder.decode()
@@ -210,4 +215,5 @@ def buffer_updated(self, nbytes: int) -> None:
                                 "record queue is full; %d record(s) to be processed",
                                 self._dbn_queue.qsize(),
                             )
-                            self._transport.pause_reading()
+                            with self._transport_lock:
+                                self._transport.pause_reading()

tests/test_live_client.py

@@ -3,7 +3,7 @@
 import pathlib
 import platform
 from io import BytesIO
-from typing import Callable
+from typing import Callable, List
 from unittest.mock import MagicMock
 
 import databento_dbn
@@ -571,11 +571,11 @@ async def test_live_async_iteration(
         symbols="TEST",
     )
 
-    live_client.start()
+    records: List[dbn.DBNStruct] = []
 
-    records = []
-    async for record in live_client:
-        records.append(record)
+    live_client.start()
+    live_client.add_callback(records.append)
+    await live_client.wait_for_close()
 
     assert len(records) == 3
     assert isinstance(records[0], databento_dbn.Metadata)