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