fix: add tests and cleanup after exceptions

pyproject.toml

@@ -41,6 +41,7 @@ test = [
     "pytest-qt",
     "pytest-asyncio",
     "pytest-timeout",
+    "fakeredis[lua]",
     "diff-cover",
     "ruff>=0.15"
 ]

src/aare/daq/auth.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import grp
 import logging
 import os

src/aare/daq/config.py

@@ -100,8 +100,8 @@ class BeamlineConfig:
             host = "localhost"
         else:
             host = cfg_get("daq.hardware.redis_url", f"{self._bl}-redis.psi.ch")
-        self.hw_lock = RedisLock(cfg.redis, name=f"{bl}:hardware_busy_lock")
         self.redis = Redis(host=host, port=6379, db=0, decode_responses=True)
+        self.hw_lock = RedisLock(self.redis, name=f"{bl}:hardware_busy_lock")
         self.simulated_detector = bl is MXBeamline.SIMULATED
         self._initialize_optional_yaml_defaults()
 
@@ -683,7 +683,7 @@ class BeamlineConfig:
     def reference_tools(self, data: SampleShortInfoList):
         self.redis.set(f"{self._bl}:reference-tools", data.model_dump_json())
 
-    def listen_changes_reference_tools(self) -> redis.client.PubSub:
+    def listen_changes_reference_tools(self) -> PubSub:
         self.redis.config_set("notify-keyspace-events", "KEA")
         pubsub = self.redis.pubsub()
         pubsub.psubscribe(f"__keyspace@0__:{self._bl}:reference-tools")
@@ -711,7 +711,7 @@ class BeamlineConfig:
         if tmp is None:
             return BeamMarkCoeffModel()
 
-        data_dict = json.loads(tmp)
+        data_dict = json.loads(tmp)  # pyright: ignore[reportArgumentType] # using sync client - remove ignore on upgrade to redis v8
         return BeamMarkCoeffModel(**data_dict)
 
     @beam_mark_coeff.setter

src/aare/daq/server.py

@@ -95,8 +95,10 @@ async def _lock_hw():
         if not hardware_busy_lock.acquire(blocking=False):
             raise BeamlineBusyException("Beamline hardware lock is held by another worker")
         logger.debug("Hardware lock acquired by process")
-        yield
-        hardware_busy_lock.release()
+        try:
+            yield
+        finally:
+            hardware_busy_lock.release()
         logger.debug("Hardware lock released by process")
 
 

tests/unit/daq/test_server_hw_lock.py

@@ -0,0 +1,73 @@
+from unittest.mock import patch
+
+import fakeredis
+import pytest
+from redis.lock import Lock as RedisLock
+
+LOCK_NAME = "SIMULATED:hardware_busy_lock"
+AUTH = {"Authorization": "Bearer fake-token"}
+
+THREAD_LOCAL = False
+
+
+@pytest.fixture
+def fake_redis():
+    return fakeredis.FakeRedis(decode_responses=True)
+
+
+@pytest.fixture
+def hw_lock(server_module, client, fake_redis, monkeypatch):
+    """This worker's lock.
+
+    Depends on ``client`` so that TestClient's lifespan (which overwrites
+    ``hardware_busy_lock``) has already run by the time we patch.
+    """
+    lock = RedisLock(fake_redis, name=LOCK_NAME, thread_local=THREAD_LOCAL)
+    monkeypatch.setattr(server_module, "hardware_busy_lock", lock)
+    return lock
+
+
+@pytest.fixture
+def other_worker(fake_redis):
+    """A second uvicorn worker's view of the same lock in Redis."""
+    return RedisLock(fake_redis, name=LOCK_NAME, thread_local=THREAD_LOCAL)
+
+
+def test_endpoint_acquires_and_releases_the_lock(client, hw_lock, other_worker):
+    with patch("aare.daq.auth.check_jwt_rw"), patch("aare.daq.server.daq"):
+        response = client.put("/beamline/omega?val=10.5", headers=AUTH)
+
+    assert response.status_code == 200
+    # Released on the way out, so a second worker can take it.
+    assert not hw_lock.owned()
+    assert other_worker.acquire(blocking=False)
+
+
+def test_endpoint_is_503_while_another_worker_holds_the_lock(client, hw_lock, other_worker):
+    assert other_worker.acquire(blocking=False)
+
+    with patch("aare.daq.auth.check_jwt_rw"), patch("aare.daq.server.daq"):
+        response = client.put("/beamline/omega?val=10.5", headers=AUTH)
+
+    assert response.status_code == 503
+    assert response.json()["exception_class"] == "BeamlineBusyException"
+    # The rejected request must not have stolen or released the other worker's lock.
+    assert other_worker.owned()
+    assert not hw_lock.owned()
+
+
+async def test_lock_hw_is_reentrant_for_the_owning_worker(server_module, hw_lock):
+    assert hw_lock.acquire(blocking=False)
+
+    async with server_module._lock_hw():
+        assert hw_lock.owned()
+
+    assert hw_lock.owned(), "re-entrant exit released a lock it did not acquire"
+
+
+async def test_lock_is_released_when_the_handler_raises(server_module, hw_lock):
+    with pytest.raises(RuntimeError):
+        async with server_module._lock_hw():
+            raise RuntimeError("handler blew up")
+
+    assert not hw_lock.owned(), "hardware lock leaked after an exception"