fix: update import for MockHandler and enhance ScienceLab mock handling

Author: minhpham1810

pslab/connection/__init__.py

@@ -5,7 +5,7 @@
 from .connection import ConnectionHandler
 from ._serial import SerialHandler
 from .wlan import WLANHandler
-from pslab.connection.mock import MockHandler
+from .mock import MockHandler
 
 __all__ = [
     "ConnectionHandler",

pslab/connection/mock.py

@@ -18,6 +18,15 @@ class MockHandler(ConnectionHandler):
     """
 
     def __init__(self, version: str = "PSLab V6 ", fw=(3, 0, 0)) -> None:
+        """Initialize MockHandler.
+
+        Parameters
+        ----------
+        version : str, optional
+            Version string to return when queried. Default is "PSLab V6 ".
+        fw : tuple of int, optional
+            Firmware version as a (major, minor, patch) tuple. Default is (3, 0, 0).
+        """
         self._rx = deque()  # bytes to be read
         self._tx = bytearray()  # bytes written by client
         self.version = version  # convenient attribute for callers
@@ -89,7 +98,7 @@ def _maybe_respond(self) -> None:
         """
         # Detect “CP.COMMON, <cmd>” patterns
         while len(self._tx) >= 2:
-            if self._tx[0] != CP.COMMON:
+            if self._tx[0] != CP.COMMON[0]:
                 # Drop unknown leading bytes
                 self._tx.pop(0)
                 continue
@@ -98,13 +107,13 @@ def _maybe_respond(self) -> None:
 
             # GET_VERSION: ConnectionHandler.get_version reads 9 bytes
             #  and checks b"PSLab"
-            if cmd == CP.GET_VERSION:
+            if cmd == CP.GET_VERSION[0]:
                 self._tx = self._tx[2:]
                 self._queue(self.version.encode("utf-8")[:9].ljust(9, b" "))
                 continue
 
             # GET_FW_VERSION: reads 3 bytes (major, minor, patch)
-            if cmd == CP.GET_FW_VERSION:
+            if cmd == CP.GET_FW_VERSION[0]:
                 self._tx = self._tx[2:]
                 major, minor, patch = self._fw
                 self._queue(bytes([major, minor, patch]))

pslab/sciencelab.py

@@ -22,16 +22,32 @@
 class ScienceLab:
     """Aggregate interface for the PSLab's instruments.
 
+    Parameters
+    ----------
+    device : ConnectionHandler, optional
+        Connection handler for communicating with the PSLab device. If not
+        provided, a new one will be created via autoconnect. If both *device*
+        and *mock* are provided, *device* takes precedence and *mock* is
+        ignored.
+    mock : bool, optional
+        If True, use a MockHandler instead of connecting to physical hardware.
+        Instruments will not be instantiated in mock mode. The default is
+        False.
+
     Attributes
     ----------
-    logic_analyzer : pslab.LogicAnalyzer
-    oscilloscope : pslab.Oscilloscope
-    waveform_generator : pslab.WaveformGenerator
-    pwm_generator : pslab.PWMGenerator
-    multimeter : pslab.Multimeter
-    power_supply : pslab.PowerSupply
+    logic_analyzer : pslab.LogicAnalyzer or None
+    oscilloscope : pslab.Oscilloscope or None
+    waveform_generator : pslab.WaveformGenerator or None
+    pwm_generator : pslab.PWMGenerator or None
+    multimeter : pslab.Multimeter or None
+    power_supply : pslab.PowerSupply or None
     i2c : pslab.I2CMaster
     nrf : pslab.peripherals.NRF24L01
+
+    Notes
+    -----
+    Instrument attributes are None when initialized with mock=True.
     """
 
     def __init__(self, device: ConnectionHandler | None = None, mock: bool = False):
@@ -43,9 +59,7 @@ def __init__(self, device: ConnectionHandler | None = None, mock: bool = False):
             self.device = autoconnect()
         self.firmware = self.device.get_firmware_version()
 
-        # mock mode initializes without hardware; instruments are not initialized.
-
-        if not mock:
+        if not mock:  # In mock mode, skip instrument initialization to avoid hardware dependencies
             self.logic_analyzer = LogicAnalyzer(device=self.device)
             self.oscilloscope = Oscilloscope(device=self.device)
             self.waveform_generator = WaveformGenerator(device=self.device)

tests/test_sciencelab_mock.py

@@ -11,10 +11,10 @@ def test_sciencelab_mock_does_not_autoconnect():
     ):
         psl = ScienceLab(mock=True)
 
-    # It should initialize and provide a firmware version object.
-    assert psl.firmware.major >= 0
-    assert psl.firmware.minor >= 0
-    assert psl.firmware.patch >= 0
+    # It should initialize and provide the expected mock firmware version object.
+    assert psl.firmware.major == 3
+    assert psl.firmware.minor == 0
+    assert psl.firmware.patch == 0
 
     # In mock mode, instruments should not be instantiated (no hardware required).
     assert psl.logic_analyzer is None