AareDAQ/tests/unit/gui/test_automation_progress_parser.py

from aare.common.automation_models import StepStatus, WorkflowStateKind
from aare.gui.threads.daq_worker import DAQWorker


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_CENTRE.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_CENTRE
    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_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_centre",'
        '"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_CENTRE
    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()