Introduce API for granular (Sensor)Reader operation

src/pms/__init__.py

@@ -21,13 +21,19 @@ class WrongMessageChecksum(SensorWarning):
     pass
 
 
-class SensorWarmingUp(SensorWarning):
+class SensorNotReady(SensorWarning):
+    """Sensor not ready to read: observations not reliable"""
+
+    pass
+
+
+class SensorWarmingUp(SensorNotReady):
     """Empty message: throw away observation and wait until sensor warms up"""
 
     pass
 
 
-class InconsistentObservation(SensorWarning):
+class InconsistentObservation(SensorNotReady):
     """Message payload makes no sense: throw away observation"""
 
     pass

src/pms/core/reader.py

@@ -18,7 +18,7 @@
 from serial import Serial
 from typer import progressbar
 
-from pms import InconsistentObservation, SensorWarmingUp, SensorWarning, logger
+from pms import SensorNotReady, SensorWarning, logger
 from pms.core import Sensor, Supported
 
 from .types import ObsData
@@ -50,7 +50,7 @@ def hexdump(self, line: Optional[int] = None) -> str:
         return f"{offset:08x}: {hex}  {dump}"
 
 
-class Reader(AbstractContextManager):
+class Reader:
     @abstractmethod
     def __call__(self, *, raw: Optional[bool]) -> Iterator[Union[RawData, ObsData]]:
         """
@@ -60,6 +60,21 @@ def __call__(self, *, raw: Optional[bool]) -> Iterator[Union[RawData, ObsData]]:
         """
         ...
 
+    @abstractmethod
+    def open(self) -> None:
+        ...
+
+    @abstractmethod
+    def close(self) -> None:
+        ...
+
+    def __enter__(self):
+        self.open()
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback) -> None:
+        self.close()
+
 
 class SensorReader(Reader):
     """Read sensor messages from serial port
@@ -78,6 +93,7 @@ def __init__(
         interval: Optional[int] = None,
         samples: Optional[int] = None,
         timeout: Optional[float] = None,
+        max_retries: Optional[int] = None,
     ) -> None:
         """Configure serial port"""
         self.sensor = sensor if isinstance(sensor, Sensor) else Sensor[sensor]
@@ -86,6 +102,7 @@ def __init__(
         self.serial.port = port
         self.serial.baudrate = self.sensor.baud
         self.serial.timeout = timeout or 5  # max time to wake up sensor
+        self.max_retries = max_retries
         self.interval = interval
         self.samples = samples
         logger.debug(
@@ -119,7 +136,7 @@ def _pre_heat(self):
         # only pre-heat the firs time
         self.pre_heat = 0
 
-    def __enter__(self) -> "SensorReader":
+    def open(self) -> None:
         """Open serial port and sensor setup"""
         if not self.serial.is_open:
             logger.debug(f"open {self.serial.port}")
@@ -143,9 +160,7 @@ def __enter__(self) -> "SensorReader":
             logger.error(f"Sensor is not {self.sensor.name}")
             raise UnableToRead("Sensor failed validation")
 
-        return self
-
-    def __exit__(self, exception_type, exception_value, traceback) -> None:
+    def close(self) -> None:
         """Put sensor to sleep and close serial port"""
         logger.debug(f"sleep {self.sensor}")
         buffer = self._cmd("sleep")
@@ -156,16 +171,23 @@ def __call__(self, *, raw: Optional[bool] = None):
         """Passive mode reading at regular intervals"""
 
         sample = 0
+        failures = 0
         while self.serial.is_open:
             try:
                 buffer = self._cmd("passive_read")
 
                 try:
                     obs = self.sensor.decode(buffer)
-                except (SensorWarmingUp, InconsistentObservation) as e:
+                except SensorNotReady as e:
+                    failures += 1
+                    if self.max_retries is not None and failures > self.max_retries:
+                        raise
                     logger.debug(e)
                     time.sleep(5)
-                except SensorWarning as e:  # pragma: no cover
+                except SensorWarning as e:
+                    failures += 1
+                    if self.max_retries is not None and failures > self.max_retries:
+                        raise
                     logger.debug(e)
                     self.serial.reset_input_buffer()
                 else:
@@ -188,14 +210,13 @@ def __init__(self, path: Path, sensor: Sensor, samples: Optional[int] = None) ->
         self.sensor = sensor
         self.samples = samples
 
-    def __enter__(self) -> "MessageReader":
+    def open(self) -> None:
         logger.debug(f"open {self.path}")
         self.csv = self.path.open()
         reader = DictReader(self.csv)
         self.data = (row for row in reader if row["sensor"] == self.sensor.name)
-        return self
 
-    def __exit__(self, exception_type, exception_value, traceback) -> None:
+    def close(self) -> None:
         logger.debug(f"close {self.path}")
         self.csv.close()
 

tests/core/test_reader.py

@@ -1,5 +1,6 @@
 import pytest
 
+import pms
 from pms.core import reader
 from pms.core.sensor import Sensor
 from tests.conftest import captured_data
@@ -12,12 +13,12 @@ def __init__(self, raise_on_enter=False):
     def __call__(self):
         raise NotImplemented
 
-    def __enter__(self):
+    def open(self):
         if self.raise_on_enter:
             raise reader.UnableToRead()
         self.entered = True
 
-    def __exit__(self, *_args):
+    def close(self):
         self.exited = True
 
 
@@ -107,20 +108,47 @@ def passive_read(n):
     )
 
 
+@pytest.fixture
+def mock_sensor_temp_failure(mock_serial):
+    def passive_read(n):
+        if n == 1:
+            # first return garbage data (bad checksum)
+            return (
+                b"BM\x00\x1c"  # expected header
+                + b"\0" * 26  # payload (to total 32 bytes)
+                + b"\x00\xFF"  # checksum
+            )
+        else:
+            # then behave like the original stub again
+            return (
+                b"BM\x00\x1c"  # expected header
+                + b".........................."  # payload
+                + b"\x05W"  # checksum
+            )
+
+    mock_serial.stub(
+        name="passive_read",
+        receive_bytes=b"BM\xe2\x00\x00\x01q",
+        send_fn=passive_read,
+    )
+
+
 @pytest.fixture
 def sensor_reader_factory(monkeypatch, mock_sensor):
     def factory(
         *,
         samples=0,  # exit immediately
         interval=None,
         sensor="PMSx003",  # match with stubs
+        max_retries=None,
     ):
         sensor_reader = reader.SensorReader(
             port=mock_sensor.port,
             samples=samples,
             interval=interval,
             sensor=sensor,
             timeout=0.01,  # low to avoid hanging on failure
+            max_retries=max_retries,
         )
 
         # https://github.com/pyserial/pyserial/issues/625
@@ -204,6 +232,47 @@ def test_sensor_reader_warm_up(
     assert len(obs) == 1
 
 
+def test_sensor_reader_warm_up_exhaust_retries(
+    mock_sensor,
+    sensor_reader_factory,
+    mock_sensor_warm_up,
+):
+    sensor_reader = sensor_reader_factory(max_retries=0)
+
+    with sensor_reader as r:
+        with pytest.raises(pms.SensorWarmingUp):
+            list(r())
+
+
+def test_sensor_reader_temp_failure(
+    mock_sensor,
+    sensor_reader_factory,
+    mock_sensor_temp_failure,
+):
+    sensor_reader = sensor_reader_factory()
+
+    with sensor_reader as r:
+        obs = list(r())
+
+    # check one sample still acquired
+    assert len(obs) == 1
+
+    # check two samples were attempted
+    assert mock_sensor.stubs["passive_read"].calls == 2
+
+
+def test_sensor_reader_temp_failure_exhaust_retries(
+    mock_sensor,
+    sensor_reader_factory,
+    mock_sensor_temp_failure,
+):
+    sensor_reader = sensor_reader_factory(max_retries=0)
+
+    with sensor_reader as r:
+        with pytest.raises(pms.SensorWarning):
+            list(r())
+
+
 def test_sensor_reader_sensor_mismatch(mock_sensor, sensor_reader_factory):
     sensor_reader = sensor_reader_factory()