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"