From 2dfd4e3d845dfc0ddb673498b15f0541d6b91844 Mon Sep 17 00:00:00 2001
From: Yash Anand <164618230+YashAnand69@users.noreply.github.com>
Date: Sat, 10 Oct 2026 16:15:18 +0530
Subject: [PATCH] fix: stop ultrasonic triggers when measurement is interrupted

---
 pslab/external/hcsr04.py     |  6 ++-
 tests/test_hcsr04_cleanup.py | 86 ++++++++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_hcsr04_cleanup.py

diff --git a/pslab/external/hcsr04.py b/pslab/external/hcsr04.py
index 602173d..bd2a15c 100644
--- a/pslab/external/hcsr04.py
+++ b/pslab/external/hcsr04.py
@@ -105,8 +105,10 @@ def estimate_distance(
             duty_cycles=self._trigger_pulse_length / self._measure_period,
         )
         # Wait one extra period to make sure we don't miss the final edge.
-        time.sleep(self._measure_period * (average + 1))
-        self._pwm.set_state(**{self._trig.lower(): 0})
+        try:
+            time.sleep(self._measure_period * (average + 1))
+        finally:
+            self._pwm.set_state(**{self._trig.lower(): 0})
         (t,) = self._la.fetch_data()
         self._sanity_check(len(t), 2 * average)
         high_times = t[1::2] - t[::2]
diff --git a/tests/test_hcsr04_cleanup.py b/tests/test_hcsr04_cleanup.py
new file mode 100644
index 0000000..0a4fcbd
--- /dev/null
+++ b/tests/test_hcsr04_cleanup.py
@@ -0,0 +1,86 @@
+"""Hardware-independent tests for ultrasonic trigger cleanup."""
+
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+
+from pslab.external.hcsr04 import HCSR04
+
+
+@pytest.fixture
+def sensor(monkeypatch):
+    """Return a sensor using controlled capture and PWM interfaces."""
+    logic = Mock()
+    pwm = Mock()
+    monkeypatch.setattr("pslab.external.hcsr04.LogicAnalyzer", Mock(return_value=logic))
+    monkeypatch.setattr("pslab.external.hcsr04.PWMGenerator", Mock(return_value=pwm))
+    sensor = HCSR04(device=Mock())
+    sensor._la.fetch_data.return_value = (np.array([100, 1100, 60100, 61100]),)
+    sensor._la.get_initial_states.return_value = {"LA1": False}
+    return sensor
+
+
+@pytest.mark.parametrize("trig", ["SQ1", "SQ4"])
+@pytest.mark.parametrize("error", [KeyboardInterrupt, RuntimeError])
+def test_interrupted_wait_stops_trigger_before_propagating(
+    sensor, monkeypatch, trig, error
+):
+    sensor._trig = trig
+    failure = error("measurement interrupted")
+    wait = Mock(side_effect=failure)
+    monkeypatch.setattr("pslab.external.hcsr04.time.sleep", wait)
+
+    with pytest.raises(error) as raised:
+        sensor.estimate_distance(average=2)
+
+    assert raised.value is failure
+    sensor._pwm.generate.assert_called_once()
+    sensor._pwm.set_state.assert_called_once_with(**{trig.lower(): 0})
+    sensor._la.fetch_data.assert_not_called()
+    sensor._la.get_initial_states.assert_not_called()
+
+
+@pytest.mark.parametrize("trig", ["SQ1", "SQ4"])
+def test_success_stops_trigger_before_fetch_and_keeps_distance(
+    sensor, monkeypatch, trig
+):
+    sensor._trig = trig
+    wait = Mock()
+    monkeypatch.setattr("pslab.external.hcsr04.time.sleep", wait)
+    calls = Mock()
+    calls.attach_mock(sensor._pwm.generate, "generate")
+    calls.attach_mock(wait, "wait")
+    calls.attach_mock(sensor._pwm.set_state, "stop")
+    calls.attach_mock(sensor._la.fetch_data, "fetch")
+
+    assert sensor.estimate_distance(average=2) == pytest.approx(0.17)
+
+    assert [call[0] for call in calls.mock_calls] == [
+        "generate",
+        "wait",
+        "stop",
+        "fetch",
+    ]
+    wait.assert_called_once_with(0.18)
+    sensor._pwm.set_state.assert_called_once_with(**{trig.lower(): 0})
+    sensor._la.capture.assert_called_once_with(channels="LA1", events=4, block=False)
+    sensor._pwm.generate.assert_called_once_with(
+        channels=trig, frequency=1 / 0.06, duty_cycles=10e-6 / 0.06
+    )
+
+
+@pytest.mark.parametrize("failure", ["timeout", "high_echo"])
+def test_capture_errors_keep_trigger_stopped(sensor, monkeypatch, failure):
+    monkeypatch.setattr("pslab.external.hcsr04.time.sleep", Mock())
+    if failure == "timeout":
+        sensor._la.fetch_data.return_value = (np.array([100, 1100]),)
+        error = TimeoutError
+    else:
+        sensor._la.get_initial_states.return_value = {"LA1": True}
+        error = RuntimeError
+
+    with pytest.raises(error):
+        sensor.estimate_distance(average=2)
+
+    sensor._pwm.set_state.assert_called_once_with(sq1=0)