AareDAQ/tests/unit/gui/test_automation_progress_parser.py

import sys
import types

from aarecommon.models.automation import StepStatus, WorkflowStateKind

from aare.gui.threads.daq_worker import DAQWorker

jfjoch_client_module = types.ModuleType("jfjoch_client")
jfjoch_client_module.ScanResult = object
jfjoch_client_module.ScanResultImagesInner = object

jfjoch_client_models_module = types.ModuleType("jfjoch_client.models")
jfjoch_client_scan_result_module = types.ModuleType("jfjoch_client.models.scan_result")
jfjoch_client_scan_result_module.ScanResult = object

sys.modules.setdefault("jfjoch_client", jfjoch_client_module)
sys.modules.setdefault("jfjoch_client.models", jfjoch_client_models_module)
sys.modules.setdefault("jfjoch_client.models.scan_result", jfjoch_client_scan_result_module)


def test_parse_automation_progress_from_sse_payload():
    progress_payload = {
        "current_step": "Center",
        "steps": [
            {
                "step": WorkflowStateKind.MOUNT.value,
                "status": StepStatus.SUCCESS.value,
                "message": "Mount complete",
            },
            {
                "step": WorkflowStateKind.LOOP_CENTER.value,
                "status": StepStatus.RUNNING.value,
                "message": "Centering sample",
            },
            {
                "step": WorkflowStateKind.RASTER.value,
                "status": StepStatus.PENDING.value,
                "message": "",
            },
            {
                "step": WorkflowStateKind.DATA_COLLECTION.value,
                "status": StepStatus.PENDING.value,
                "message": "",
            },
            {
                "step": WorkflowStateKind.FINAL.value,
                "status": StepStatus.PENDING.value,
                "message": "",
            },
        ],
        "finished": False,
        "success": None,
    }

    progress = DAQWorker._parse_automation_progress(progress_payload)

    assert progress.current_step == "Center"
    assert progress.finished is False
    assert progress.success is None

    assert len(progress.steps) == 5

    assert progress.steps[0].step == WorkflowStateKind.MOUNT
    assert progress.steps[0].status == StepStatus.SUCCESS
    assert progress.steps[0].message == "Mount complete"

    assert progress.steps[1].step == WorkflowStateKind.LOOP_CENTER
    assert progress.steps[1].status == StepStatus.RUNNING
    assert progress.steps[1].message == "Centering sample"

    assert progress.steps[2].step == WorkflowStateKind.RASTER
    assert progress.steps[2].status == StepStatus.PENDING
    assert progress.steps[2].message == ""


def test_parse_automation_progress_finished_state():
    progress_payload = {
        "current_step": "Paused/Finished",
        "steps": [
            {
                "step": WorkflowStateKind.FINAL.value,
                "status": StepStatus.SUCCESS.value,
                "message": "Automation complete",
            }
        ],
        "finished": True,
        "success": True,
    }

    progress = DAQWorker._parse_automation_progress(progress_payload)

    assert progress.current_step == "Paused/Finished"
    assert progress.finished is True
    assert progress.success is True
    assert len(progress.steps) == 1
    assert progress.steps[0].step == WorkflowStateKind.FINAL
    assert progress.steps[0].status == StepStatus.SUCCESS
    assert progress.steps[0].message == "Automation complete"


def test_parse_automation_progress_parses_events():
    progress_payload = {
        "current_step": "Center",
        "steps": [],
        "events": [
            {
                "ts": "2026-06-03T10:00:00+00:00",
                "level": "warning",
                "code": "LOOP_CENTERING_FAILED",
                "exception_class": "LoopCenteringFailed",
                "message": "Loop centering failed",
                "sample_id": 42,
                "context": {"step": "loop_center"},
            }
        ],
        "finished": False,
        "success": None,
    }

    progress = DAQWorker._parse_automation_progress(progress_payload)

    assert len(progress.events) == 1
    assert progress.events[0].level == "WARNING"
    assert progress.events[0].code == "LOOP_CENTERING_FAILED"
    assert progress.events[0].sample_id == 42


def test_handle_automation_progress_event_dedups_events_by_timestamp(caplog):
    worker = DAQWorker(base_url=None, token="test-token")

    payload = (
        '{"seq":1,"progress":{"current_step":"Center","steps":[],"events":['
        '{"ts":"2026-06-03T10:00:00+00:00","level":"WARNING","code":"LOOP_CENTERING_FAILED",'
        '"exception_class":"LoopCenteringFailed","message":"Loop centering failed","sample_id":7,'
        '"context":{"step":"loop_center"}}],"finished":false,"success":null}}'
    )

    caplog.clear()
    worker._handle_automation_progress_event(payload)
    worker._handle_automation_progress_event(payload)

    warning_messages = [
        record.message for record in caplog.records if record.levelname == "WARNING"
    ]
    assert warning_messages.count("Loop centering failed") == 1

    worker.cleanup()


def test_handle_automation_progress_event_trips_recurrence_watcher():
    worker = DAQWorker(base_url=None, token="test-token")
    worker._recurrence_watchers = [w for w in worker._recurrence_watchers if w.name == "alc"]

    trips: list[str] = []
    worker.automation_critical_failure.connect(trips.append)

    payload = (
        '{"seq":1,"progress":{"current_step":"Center","steps":[],"events":['
        '{"ts":"2026-06-03T10:00:00+00:00","level":"WARNING","code":"LOOP_CENTERING_FAILED",'
        '"exception_class":"LoopCenteringFailed","message":"Loop centering failed","sample_id":7,'
        '"context":{"step":"loop_center"}}],"finished":false,"success":null}}'
    )

    worker._handle_automation_progress_event(payload)
    worker._handle_automation_progress_event(
        payload.replace('"seq":1', '"seq":2').replace("10:00:00", "10:00:01")
    )
    worker._handle_automation_progress_event(
        payload.replace('"seq":1', '"seq":3').replace("10:00:00", "10:00:02")
    )

    assert len(trips) == 1
    assert "consecutive alc errors" in trips[0]

    worker.cleanup()


def test_process_automation_progress_buffer_waits_for_complete_event():
    worker = DAQWorker(base_url=None, token="test-token")
    emitted: list = []

    worker.automation_progress.connect(emitted.append)

    worker._automation_progress_buffer = (
        'data: {"seq":1,"progress":{"current_step":"Center","steps":[{"step":"mount",'
        '"status":"success","message":"Mount complete"}],"finished":false,"success":null}}'
    )

    worker._process_automation_progress_buffer()

    assert emitted == []
    assert worker._automation_progress_buffer.startswith("data:")

    worker.cleanup()


def test_process_automation_progress_buffer_handles_split_sse_event():
    worker = DAQWorker(base_url=None, token="test-token")
    emitted: list = []

    worker.automation_progress.connect(emitted.append)

    first_chunk = 'data: {"seq":1,"progress":{"current_step":"Center","steps":[{"step":"mount",'
    second_chunk = (
        '"status":"success","message":"Mount complete"},{"step":"loop_center",'
        '"status":"running","message":"Centering sample"}],'
        '"finished":false,"success":null}}\n\n'
    )

    worker._automation_progress_buffer += first_chunk
    worker._process_automation_progress_buffer()

    assert emitted == []

    worker._automation_progress_buffer += second_chunk
    worker._process_automation_progress_buffer()

    assert len(emitted) == 1

    progress = emitted[0]
    assert progress.current_step == "Center"
    assert progress.finished is False
    assert progress.success is None
    assert len(progress.steps) == 2
    assert progress.steps[0].step == WorkflowStateKind.MOUNT
    assert progress.steps[0].status == StepStatus.SUCCESS
    assert progress.steps[1].step == WorkflowStateKind.LOOP_CENTER
    assert progress.steps[1].status == StepStatus.RUNNING
    assert worker._automation_progress_buffer == ""

    worker.cleanup()


def test_process_automation_progress_buffer_handles_multiple_events():
    worker = DAQWorker(base_url=None, token="test-token")
    emitted: list = []

    worker.automation_progress.connect(emitted.append)

    worker._automation_progress_buffer = (
        'data: {"seq":1,"progress":{"current_step":"Mount","steps":[{"step":"mount",'
        '"status":"running","message":"Mounting sample"}],"finished":false,"success":null}}\n\n'
        'data: {"seq":2,"progress":{"current_step":"Paused/Finished","steps":[{"step":"final",'
        '"status":"success","message":"Automation complete"}],"finished":true,"success":true}}\n\n'
    )

    worker._process_automation_progress_buffer()

    assert len(emitted) == 2

    first = emitted[0]
    assert first.current_step == "Mount"
    assert first.finished is False
    assert first.steps[0].step == WorkflowStateKind.MOUNT
    assert first.steps[0].status == StepStatus.RUNNING

    second = emitted[1]
    assert second.current_step == "Paused/Finished"
    assert second.finished is True
    assert second.success is True
    assert second.steps[0].step == WorkflowStateKind.FINAL
    assert second.steps[0].status == StepStatus.SUCCESS

    assert worker._automation_progress_buffer == ""

    worker.cleanup()


def test_handle_automation_progress_event_ignores_null_progress():
    worker = DAQWorker(base_url=None, token="test-token")
    emitted: list = []

    worker.automation_progress.connect(emitted.append)

    worker._handle_automation_progress_event('{"seq":3,"progress":null}')

    assert emitted == []

    worker.cleanup()