AareDAQ/tests/unit/daq/test_tellupdater.py

import json
from unittest.mock import MagicMock, patch

from aare.daq import tellupdater


def test_compare_and_report_change():
    class MockPuck:
        def __init__(self, puck_name, dewar_id, pos):
            self.puck_name = puck_name
            self.dewar_id = dewar_id
            self.tell_position = pos

    old = [MockPuck("p1", 1, "A1"), MockPuck("p2", 1, "A2")]
    new = [MockPuck("p2", 1, "A2"), MockPuck("p3", 1, "A3")]

    joined, left = tellupdater.compare_and_report_change(old, new)
    assert joined == {("p3", 1, "A3")}
    assert left == {("p1", 1, "A1")}


def test_compare_and_report_change_ignore_x1():
    class MockPuck:
        def __init__(self, puck_name, dewar_id, pos):
            self.puck_name = puck_name
            self.dewar_id = dewar_id
            self.tell_position = pos

    old = [MockPuck("p1", 1, "X1")]
    new = [MockPuck("p2", 1, "X1")]

    joined, left = tellupdater.compare_and_report_change(old, new)
    assert joined == set()
    assert left == set()


@patch("aare.daq.tellupdater.tell_client")
@patch("aare.daq.tellupdater.aare_db")
def test_handle_tell_change_event(mock_db, mock_tell):
    mock_tell.get_detected_pucks.return_value = ["puck1", "puck2"]
    tellupdater.handle_tell_change_event()
    mock_tell.get_detected_pucks.assert_called_once()
    mock_db.set_pucks_beamline.assert_called_with(["puck1", "puck2"])


@patch("aare.daq.tellupdater.tell_client")
def test_ws_update_samples_info(mock_tell):
    pucks = [{"id": "1"}]
    tellupdater.ws_update_samples_info(pucks)
    mock_tell.set_samples_info.assert_called_with(pucks)


def test_on_sse_event():
    class MockEvent:
        def __init__(self, event, data):
            self.event = event
            self.data = data

    event = MockEvent("DewarContentUpdate", "some data")
    with patch("aare.daq.tellupdater.handle_tell_change_event") as mock_handle:
        tellupdater.on_sse_event(event)
        mock_handle.assert_called_once()


def test_extract_tracked_tell_event_with_json_payload():
    tracked = tellupdater.extract_tracked_tell_event(
        "Gripper detection", json.dumps({"value": "Pin still in Gripper"})
    )
    assert tracked == ("Gripper detection", "Pin still in Gripper")


def test_extract_tracked_tell_event_for_standalone_message():
    tracked = tellupdater.extract_tracked_tell_event(
        "Motion Sync", json.dumps({"message": "Sample put on Gonio"})
    )
    assert tracked == ("Motion Sync", "Sample put on Gonio")


def test_on_sse_event_tracks_interesting_event():
    class MockEvent:
        def __init__(self, event, data):
            self.event = event
            self.data = data

    tellupdater.latest_tell_events.clear()
    event = MockEvent("Motion Task", json.dumps({"value": "dry"}))

    with patch("aare.daq.tellupdater.handle_tell_change_event") as mock_handle:
        tellupdater.on_sse_event(event)
        mock_handle.assert_not_called()

    assert tellupdater.latest_tell_events["Motion Task"] == "dry"


def test_record_tell_event_logs_to_journal():
    tellupdater.latest_tell_events.clear()
    tellupdater.tell_event_history.clear()

    with patch("aare.daq.tellupdater.logger.info") as mock_info:
        tellupdater.record_tell_event("Motion Sync", "Sample put on Gonio")

    assert tellupdater.latest_tell_events["Motion Sync"] == "Sample put on Gonio"
    mock_info.assert_called_once_with("[TELL][JOURNAL] event=Motion Sync value=Sample put on Gonio")


def test_record_tell_event_writes_history_to_redis():
    tellupdater.latest_tell_events.clear()
    tellupdater.tell_event_history.clear()

    mock_config = MagicMock()
    mock_config._bl = "x10sa"

    with patch("aare.daq.tellupdater.config", mock_config):
        tellupdater.record_tell_event("Motion Task", "dry")

    mock_config._client.set.assert_called_once()
    redis_key, redis_value = mock_config._client.set.call_args.args
    assert redis_key == "x10sa:tell_events"

    payload = json.loads(redis_value)
    assert len(payload) == 1
    assert payload[0]["class"] == "Motion Task"
    assert payload[0]["event"] == "dry"
    assert "timestamp" in payload[0]


def test_record_tell_event_keeps_last_25_events():
    tellupdater.latest_tell_events.clear()
    tellupdater.tell_event_history.clear()

    mock_config = MagicMock()
    mock_config._bl = "x10sa"

    with patch("aare.daq.tellupdater.config", mock_config):
        for idx in range(30):
            tellupdater.record_tell_event("Motion Sync", f"event-{idx}")

    redis_key, redis_value = mock_config._client.set.call_args.args
    assert redis_key == "x10sa:tell_events"

    payload = json.loads(redis_value)
    assert len(payload) == 25
    assert payload[0]["event"] == "event-5"
    assert payload[-1]["event"] == "event-29"


def test_on_message_logs_json_decode_error():
    with patch("aare.daq.tellupdater.logger.exception") as mock_exception:
        tellupdater.on_message(None, "{invalid")

    mock_exception.assert_called_once_with("[WS] Failed to process message")


@patch("aare.daq.tellupdater.tell_client")
@patch("aare.daq.tellupdater.aare_db")
def test_handle_tell_change_event_logs_failure(mock_db, mock_tell):
    mock_tell.get_detected_pucks.side_effect = ValueError("boom")

    with patch("aare.daq.tellupdater.logger.exception") as mock_exception:
        tellupdater.handle_tell_change_event()

    mock_exception.assert_called_once_with("[SSE] Failed to update puck state after TELL change")


@patch("aare.daq.tellupdater.websocket.WebSocketApp")
def test_on_message(mock_ws_app):
    message = json.dumps(
        [
            {
                "id": 1,
                "barcode": "B1",
                "position": "P1",
                "puck_name": "Puck1",
                "puck_type": "Type1",
                "puck_location_in_dewar": 1,
                "dewar_id": 1,
                "dewar_name": "Dewar1",
                "pgroup": "p12345",
                "tell_position": "A1",
            }
        ]
    )
    with (
        patch("aare.daq.tellupdater.ws_update_samples_info") as mock_update,
        patch("aare.daq.tellupdater.handle_tell_change_event"),
    ):
        tellupdater.current_pucks = []
        tellupdater.on_message(None, message)
        mock_update.assert_called_once()