FIX: Live test flakiness

Author: nmacholl

databento/live/client.py

@@ -121,14 +121,14 @@ async def __anext__(self) -> DBNRecord:
 
     def __iter__(self) -> "Live":
         logger.debug("starting iteration")
-        self._dbn_queue._enabled = True
+        self._dbn_queue._enabled.set()
         return self
 
     def __next__(self) -> DBNRecord:
         if self._dbn_queue is None:
             raise ValueError("iteration has not started")
 
-        while not self._session.is_disconnected() or self._dbn_queue.qsize() > 0:
+        while not self._session.is_disconnected() or self._dbn_queue._qsize() > 0:
             try:
                 record = self._dbn_queue.get(block=False)
             except queue.Empty:
@@ -141,14 +141,14 @@ def __next__(self) -> DBNRecord:
                 self._dbn_queue.task_done()
                 return record
             finally:
-                if not self._dbn_queue.full() and not self._session.is_reading():
+                if not self._dbn_queue.half_full() and not self._session.is_reading():
                     logger.debug(
                         "resuming reading with %d pending records",
-                        self._dbn_queue.qsize(),
+                        self._dbn_queue._qsize(),
                     )
                     self._session.resume_reading()
 
-        self._dbn_queue._enabled = False
+        self._dbn_queue._enabled.clear()
         raise StopIteration
 
     def __repr__(self) -> str:

databento/live/session.py

@@ -3,6 +3,7 @@
 import logging
 import queue
 import struct
+import threading
 from typing import IO, Callable, Iterable, List, Optional, Set, Union
 
 import databento_dbn
@@ -38,40 +39,23 @@ class DBNQueue(queue.Queue):  # type: ignore [type-arg]
 
     def __init__(self, maxsize: int) -> None:
         super().__init__(maxsize)
-        self._enabled: bool = False
-
-    def put(
-        self,
-        item: DBNRecord,
-        block: bool = True,
-        timeout: Optional[float] = None,
-    ) -> None:
-        if self._enabled:
-            return super().put(item, block, timeout)
-        raise ValueError("queue is not being iterated")
-
-    def put_nowait(
-        self,
-        item: DBNRecord,
-    ) -> None:
-        if self._enabled:
-            return super().put_nowait(item)
-        raise ValueError("queue is not being iterated")
+        self._enabled = threading.Event()
 
     @property
     def enabled(self) -> bool:
         """
         True if the Queue will allow pushing.
         A queue should only be enabled when it has a consumer.
         """
-        return self._enabled
+        return self._enabled.is_set()
 
-    def full(self) -> bool:
+    def half_full(self) -> bool:
         """
         Implementation which reports the queue as full when it
         has reached half capacity.
         """
-        return self.qsize() > self.maxsize // 2
+        with self.mutex:
+            return self._qsize() > self.maxsize // 2
 
 
 @dataclasses.dataclass
@@ -180,10 +164,10 @@ def received_record(self, record: DBNRecord) -> None:
                     record.ts_event,
                 )
             else:
-                if self._dbn_queue.full():
+                if self._dbn_queue.half_full():
                     logger.warning(
                         "record queue is full; %d record(s) to be processed",
-                        self._dbn_queue.qsize(),
+                        self._dbn_queue._qsize(),
                     )
                     self.transport.pause_reading()
 

tests/data/generator.py

@@ -11,14 +11,21 @@
         print(schema.value)
         path = f"test_data.{schema.value}.dbn.zst"
 
+        if schema == Schema.IMBALANCE:
+            dataset = "XNAS.ITCH"
+            symbol = "NVDA"
+        else:
+            dataset = "GLBX.MDP3"
+            symbol = "ESH1"
+
         # Execute request through client
         data: DBNStore = client.timeseries.get_range(
-            dataset="GLBX.MDP3",
-            symbols=["ESH1"],
+            dataset=dataset,
+            symbols=[symbol],
             schema=schema,
             start="2020-12-28T13:00",
             end="2020-12-29T13:00",
-            limit=2,  # <-- limiting response to 2 records only (for test cases)
+            limit=4,  # <-- limiting response to 4 records only (for test cases)
             path=path,
         )  # -> DBNStore