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()

tests/unit/queue/test_queue_models.py

@@ -0,0 +1,131 @@
+import json
+from uuid import UUID, uuid4
+
+import pytest
+
+from aare.queue.models import (
+    QueueInstructionBase,
+    QueueItemKind,
+    RasterInstruction,
+    RotationInstruction,
+    ScreeningInstruction,
+    UnknownInstruction,
+    dump_queue_item,
+    parse_queue_item,
+)
+
+KNOWN_INSTRUCTIONS = [
+    (RasterInstruction, QueueItemKind.RASTER),
+    (ScreeningInstruction, QueueItemKind.SCREENING),
+    (RotationInstruction, QueueItemKind.ROTATION),
+]
+
+# A plausible instruction from a producer newer than this reader: a kind we have
+# never heard of, plus fields we know nothing about.
+FUTURE_INSTRUCTION = {
+    "id": "11111111-2222-3333-4444-555555555555",
+    "kind": "xrf_map",
+    "sample_id": 9,
+    "grid_steps": [4, 4],
+    "dwell_ms": 50,
+}
+
+
+@pytest.mark.parametrize("cls,kind", KNOWN_INSTRUCTIONS, ids=lambda v: getattr(v, "value", None))
+def test_known_instruction_defaults_its_own_kind(cls, kind):
+    assert cls(sample_id=1).kind == kind
+
+
+def test_ids_are_generated_and_distinct():
+    ids = {RasterInstruction(sample_id=1).id for _ in range(5)}
+    assert len(ids) == 5
+
+
+@pytest.mark.parametrize("cls,kind", KNOWN_INSTRUCTIONS, ids=lambda v: getattr(v, "value", None))
+def test_known_instruction_round_trips_to_its_own_class(cls, kind):
+    original = cls(sample_id=3)
+    restored = parse_queue_item(dump_queue_item(original))
+    assert type(restored) is cls
+    assert restored.sample_id == 3
+
+
+@pytest.mark.parametrize("cls,kind", KNOWN_INSTRUCTIONS, ids=lambda v: getattr(v, "value", None))
+def test_round_trip_preserves_the_id(cls, kind):
+    """Identity ops target the id, so a read must never regenerate it."""
+    original = cls(sample_id=3)
+    assert parse_queue_item(dump_queue_item(original)).id == original.id
+
+
+def test_unrecognised_kind_becomes_an_unknown_instruction():
+    item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION))
+    assert isinstance(item, UnknownInstruction)
+    assert item.kind == "xrf_map"
+    assert item.id == UUID(FUTURE_INSTRUCTION["id"])
+
+
+def test_unknown_instruction_survives_a_read_write_cycle_verbatim():
+    """Reordering rewrites every entry. A reader that cannot interpret an item
+    must not drop the fields it did not understand."""
+    item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION))
+    assert json.loads(dump_queue_item(item)) == FUTURE_INSTRUCTION
+
+
+def test_unknown_instruction_is_a_queue_instruction():
+    """It shares the base, so consumers reach its common fields the same way
+    they would for a kind they understand."""
+    item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION))
+    assert isinstance(item, QueueInstructionBase)
+    assert item.id == UUID(FUTURE_INSTRUCTION["id"])
+    assert item.sample_id == FUTURE_INSTRUCTION["sample_id"]
+
+
+def test_unknown_fields_are_kept_as_extras():
+    """Held as pydantic extras rather than a raw blob, which is what lets them
+    serialize back out without a custom serializer."""
+    item = parse_queue_item(json.dumps(FUTURE_INSTRUCTION))
+    assert item.model_extra == {
+        "grid_steps": FUTURE_INSTRUCTION["grid_steps"],
+        "dwell_ms": FUTURE_INSTRUCTION["dwell_ms"],
+    }
+
+
+def test_unknown_instruction_does_not_regenerate_a_missing_id():
+    """The base defaults id, but a stored entry always has one already --
+    minting a fresh one on read would silently break remove/reorder."""
+    with pytest.raises(ValueError):
+        parse_queue_item(json.dumps({"kind": "xrf_map", "sample_id": 9}))
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        {"kind": "xrf_map", "sample_id": 9},
+        {"id": "not-a-uuid", "kind": "xrf_map", "sample_id": 9},
+        {"id": str(uuid4()), "sample_id": 9},
+        {"id": str(uuid4()), "kind": "xrf_map"},
+    ],
+    ids=["no-id", "non-uuid-id", "no-kind", "no-sample-id"],
+)
+def test_entries_missing_a_base_field_are_rejected(payload):
+    """Every instruction inherits id and sample_id and declares a kind, so an
+    entry missing any of them is corrupt rather than merely unrecognised -- it
+    must not parse as Unknown."""
+    with pytest.raises(ValueError):
+        parse_queue_item(json.dumps(payload))
+
+
+def test_malformed_json_is_rejected():
+    with pytest.raises(ValueError):
+        parse_queue_item("NOT JSON AT ALL")
+
+
+def test_non_object_json_is_rejected():
+    with pytest.raises(TypeError):
+        parse_queue_item("[1, 2, 3]")
+
+
+def test_a_known_kind_with_a_bad_payload_is_not_smuggled_in_as_unknown():
+    """Falling back on any ValidationError would turn a malformed raster item
+    into an opaque Unknown that silently round-trips broken data."""
+    with pytest.raises(ValueError):
+        parse_queue_item(json.dumps({"id": str(uuid4()), "kind": "raster"}))  # no sample_id