wip feat: add redis-based queue

tests/unit/queue/test_queue_client.py

@@ -0,0 +1,287 @@
+import gc
+import json
+import threading
+from uuid import UUID, uuid4
+
+import fakeredis
+import pytest
+from aarecommon.models.beamline import MXBeamline
+
+from aare.queue.client import AareQueueClient
+from aare.queue.models import (
+    QUEUE_KEY_SUFFIX,
+    QUEUE_NOTIFICATION_CHANNEL_SUFFIX,
+    RasterInstruction,
+    RotationInstruction,
+    UnknownInstruction,
+    dump_queue_item,
+)
+
+BEAMLINE = MXBeamline.X10SA
+QUEUE_KEY = str(BEAMLINE) + QUEUE_KEY_SUFFIX
+NOTIFY_CHANNEL = str(BEAMLINE) + QUEUE_NOTIFICATION_CHANNEL_SUFFIX
+
+# Callbacks run on the client's dispatch thread, so tests wait on an Event
+# rather than sleeping. Generous enough for a loaded CI box, but only ever
+# paid in full when something is actually broken.
+DISPATCH_TIMEOUT_S = 2.0
+# Waiting for something NOT to happen has to be a real sleep; keep it short.
+QUIET_PERIOD_S = 0.3
+
+
+class Recorder:
+    """A callback that records its calls, as a bound method.
+
+    Bound methods are the realistic case (a Qt panel registering
+    ``self.on_queue_update``) and the one that a weak callback registry drops.
+    """
+
+    def __init__(self):
+        self.calls = 0
+        self.called = threading.Event()
+
+    def on_queue_update(self):
+        self.calls += 1
+        self.called.set()
+
+    def wait(self, timeout=DISPATCH_TIMEOUT_S) -> bool:
+        return self.called.wait(timeout)
+
+
+@pytest.fixture
+def redis():
+    return fakeredis.FakeRedis(decode_responses=True)
+
+
+@pytest.fixture
+def client(redis):
+    c = AareQueueClient(BEAMLINE, redis)
+    yield c
+    c.shutdown()
+
+
+@pytest.fixture
+def recorder(client):
+    r = Recorder()
+    client.add_update_callback(r.on_queue_update)
+    return r
+
+
+def notify(redis):
+    """Stand in for whatever writes the queue and announces the change."""
+    redis.publish(NOTIFY_CHANNEL, "1")
+
+
+# ---------------------------------------------------------------------------
+# construction and notification plumbing
+# ---------------------------------------------------------------------------
+
+
+def test_client_constructs_with_a_subscribed_handler(client):
+    """Handlers must be passed to subscribe() as kwargs. A dict passed
+    positionally is read as a list of channel names, leaving the channel
+    handler-less and making run_in_thread raise PubSubError."""
+    assert NOTIFY_CHANNEL in client._notification_sub.channels
+    assert client._notification_sub.channels[NOTIFY_CHANNEL] is not None
+
+
+def test_callback_runs_when_a_notification_arrives(redis, recorder):
+    notify(redis)
+    assert recorder.wait(), "callback never ran"
+    assert recorder.calls == 1
+
+
+def test_callback_does_not_rerun_without_a_new_notification(redis, recorder):
+    """The dispatch event must be cleared after waking. Left set, the loop
+    spins and replays every callback continuously."""
+    notify(redis)
+    assert recorder.wait()
+
+    calls_after_first = recorder.calls
+    threading.Event().wait(QUIET_PERIOD_S)
+    assert recorder.calls == calls_after_first
+
+
+def test_each_notification_dispatches_again(redis, recorder):
+    for expected in (1, 2, 3):
+        recorder.called.clear()
+        notify(redis)
+        assert recorder.wait(), f"no dispatch for notification {expected}"
+        assert recorder.calls == expected
+
+
+# ---------------------------------------------------------------------------
+# callback registry
+# ---------------------------------------------------------------------------
+
+
+def test_bound_method_callback_survives_registration(client):
+    """A bound method is a fresh object per attribute access, so a weak
+    registry collects it the instant registration returns."""
+    recorder = Recorder()
+    cb_id = client.add_update_callback(recorder.on_queue_update)
+    gc.collect()
+    assert cb_id in client._callbacks
+
+
+def test_lambda_callback_survives_registration(client):
+    cb_id = client.add_update_callback(lambda: None)
+    gc.collect()
+    assert cb_id in client._callbacks
+
+
+def test_every_registered_callback_is_dispatched(redis, client):
+    recorders = [Recorder() for _ in range(3)]
+    for r in recorders:
+        client.add_update_callback(r.on_queue_update)
+
+    notify(redis)
+
+    for i, r in enumerate(recorders):
+        assert r.wait(), f"callback {i} never ran"
+
+
+def test_removed_callback_stops_being_dispatched(redis, client):
+    removed, kept = Recorder(), Recorder()
+    removed_id = client.add_update_callback(removed.on_queue_update)
+    client.add_update_callback(kept.on_queue_update)
+
+    client.remove_update_callback(removed_id)
+    notify(redis)
+
+    assert kept.wait(), "surviving callback never ran"
+    assert removed.calls == 0
+
+
+def test_removing_an_unregistered_callback_raises(client):
+    with pytest.raises(ValueError):
+        client.remove_update_callback(str(uuid4()))
+
+
+def test_callback_may_deregister_itself_during_dispatch(redis, client):
+    """The dispatch loop iterates a snapshot; mutating the live dict under the
+    reentrant lock would raise RuntimeError mid-dispatch."""
+    done = threading.Event()
+    ids = {}
+
+    def self_removing():
+        client.remove_update_callback(ids["self"])
+        done.set()
+
+    ids["self"] = client.add_update_callback(self_removing)
+    later = Recorder()
+    client.add_update_callback(later.on_queue_update)
+
+    notify(redis)
+
+    assert done.wait(DISPATCH_TIMEOUT_S), "self-removing callback never ran"
+    assert later.wait(), "dispatch aborted before reaching the next callback"
+
+
+def test_a_raising_callback_does_not_stop_the_others(redis, client):
+    def boom():
+        raise RuntimeError("callback blew up")
+
+    client.add_update_callback(boom)
+    survivor = Recorder()
+    client.add_update_callback(survivor.on_queue_update)
+
+    notify(redis)
+
+    assert survivor.wait(), "one bad callback killed the dispatch"
+
+
+def test_dispatch_survives_a_raising_callback_across_notifications(redis, client):
+    """An exception must not take the dispatch thread down with it."""
+
+    def boom():
+        raise RuntimeError("callback blew up")
+
+    client.add_update_callback(boom)
+    survivor = Recorder()
+    client.add_update_callback(survivor.on_queue_update)
+
+    notify(redis)
+    assert survivor.wait()
+
+    survivor.called.clear()
+    notify(redis)
+    assert survivor.wait(), "dispatch thread died on the first exception"
+
+
+# ---------------------------------------------------------------------------
+# reads
+# ---------------------------------------------------------------------------
+
+
+def test_get_all_is_empty_for_an_untouched_queue(client):
+    assert client.get_all() == []
+
+
+def test_get_all_preserves_insertion_order(redis, client):
+    items = [RasterInstruction(sample_id=i) for i in range(4)]
+    redis.rpush(QUEUE_KEY, *[dump_queue_item(i) for i in items])
+
+    assert [i.id for i in client.get_all()] == [i.id for i in items]
+
+
+def test_get_all_keeps_unrecognised_instructions(redis, client):
+    future = {"id": str(uuid4()), "kind": "xrf_map", "sample_id": 9, "dwell_ms": 50}
+    redis.rpush(QUEUE_KEY, json.dumps(future))
+
+    (item,) = client.get_all()
+    assert isinstance(item, UnknownInstruction)
+    assert item.kind == "xrf_map"
+
+
+def test_get_all_skips_corrupt_entries_without_losing_the_rest(redis, client):
+    """One bad write must not blank the queue for every consumer."""
+    good = RotationInstruction(sample_id=1)
+    redis.rpush(
+        QUEUE_KEY,
+        dump_queue_item(good),
+        "NOT JSON AT ALL",
+        "[1, 2, 3]",
+        json.dumps({"kind": "xrf_map", "sample_id": 9}),  # no id
+        dump_queue_item(RasterInstruction(sample_id=2)),
+    )
+
+    items = client.get_all()
+
+    assert len(items) == 2
+    assert items[0].id == good.id
+
+
+def test_every_item_read_back_is_targetable_by_id(redis, client):
+    """Identity ops (remove, reorder) rely on this holding for unknown items too."""
+    redis.rpush(
+        QUEUE_KEY,
+        dump_queue_item(RasterInstruction(sample_id=1)),
+        json.dumps({"id": str(uuid4()), "kind": "xrf_map", "sample_id": 2}),
+    )
+
+    items = client.get_all()
+
+    assert len(items) == 2, "an unknown instruction was dropped, making this vacuous"
+    assert all(isinstance(item.id, UUID) for item in items)
+
+
+# ---------------------------------------------------------------------------
+# shutdown
+# ---------------------------------------------------------------------------
+
+
+def test_shutdown_stops_both_threads(redis):
+    client = AareQueueClient(BEAMLINE, redis)
+    client.shutdown()
+
+    assert not client._callback_thread.is_alive()
+    assert not client._notification_thread.is_alive()
+
+
+def test_shutdown_is_safe_with_callbacks_registered(redis):
+    client = AareQueueClient(BEAMLINE, redis)
+    client.add_update_callback(Recorder().on_queue_update)
+    client.shutdown()
+
+    assert not client._callback_thread.is_alive()