106 lines
3.9 KiB
Python
106 lines
3.9 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from aare.common.automation_workflow import (
|
|
STATE_REGISTRY, HANDLER_REGISTRY, SIMULATED_HANDLER_REGISTRY,
|
|
get_state_definition, get_allowed_next_states, can_transition,
|
|
MountHandler, LoopCentreHandler, RasterHandler, DataCollectionHandler,
|
|
WorkflowRunner, WorkflowContext, WorkflowMode, WorkflowStateKind, StepStatus
|
|
)
|
|
|
|
def test_get_state_definition():
|
|
defn = get_state_definition(WorkflowStateKind.MOUNT)
|
|
assert defn.kind == WorkflowStateKind.MOUNT
|
|
|
|
with pytest.raises(KeyError, match="Unknown workflow state"):
|
|
get_state_definition("invalid_state")
|
|
|
|
def test_get_allowed_next_states():
|
|
# Mount -> Loop Centre
|
|
next_states = get_allowed_next_states(WorkflowStateKind.MOUNT, mode=WorkflowMode.AUTOMATION)
|
|
assert WorkflowStateKind.LOOP_CENTRE in next_states
|
|
|
|
# Loop Centre -> Raster or Data Collection (depending on mode)
|
|
next_states_auto = get_allowed_next_states(WorkflowStateKind.LOOP_CENTRE, mode=WorkflowMode.AUTOMATION)
|
|
assert WorkflowStateKind.RASTER in next_states_auto
|
|
assert WorkflowStateKind.DATA_COLLECTION not in next_states_auto
|
|
|
|
next_states_manual = get_allowed_next_states(WorkflowStateKind.LOOP_CENTRE, mode=WorkflowMode.GUIDED_MANUAL)
|
|
assert WorkflowStateKind.RASTER in next_states_manual
|
|
assert WorkflowStateKind.DATA_COLLECTION in next_states_manual
|
|
|
|
def test_can_transition():
|
|
assert can_transition(WorkflowStateKind.MOUNT, WorkflowStateKind.LOOP_CENTRE) is True
|
|
assert can_transition(WorkflowStateKind.MOUNT, WorkflowStateKind.RASTER) is False
|
|
|
|
@pytest.fixture
|
|
def context():
|
|
return WorkflowContext(
|
|
mode=WorkflowMode.AUTOMATION,
|
|
queue_id="q1",
|
|
item_id="item1",
|
|
sample_id=123
|
|
)
|
|
|
|
def test_handlers_execute(context):
|
|
handlers = [
|
|
MountHandler(),
|
|
LoopCentreHandler(),
|
|
RasterHandler(),
|
|
DataCollectionHandler()
|
|
]
|
|
|
|
for handler in handlers:
|
|
res = handler.execute(context)
|
|
assert res.status == StepStatus.SUCCESS
|
|
assert context.current_state == handler.state_kind
|
|
|
|
def test_handler_abort(context):
|
|
context.abort_requested = True
|
|
handler = MountHandler()
|
|
with pytest.raises(RuntimeError, match="Abort requested"):
|
|
handler.execute(context)
|
|
|
|
def test_workflow_runner(context):
|
|
runner = WorkflowRunner()
|
|
|
|
# Start at None, go to Mount
|
|
res = runner.run_state(context, WorkflowStateKind.MOUNT)
|
|
assert res.state == WorkflowStateKind.MOUNT
|
|
assert context.current_step_index == 1
|
|
|
|
# Mount -> Loop Centre
|
|
res = runner.run_state(context, WorkflowStateKind.LOOP_CENTRE)
|
|
assert res.state == WorkflowStateKind.LOOP_CENTRE
|
|
|
|
# Invalid transition: Loop Centre -> Mount (in reverse)
|
|
with pytest.raises(RuntimeError, match="Transition not allowed"):
|
|
runner.run_state(context, WorkflowStateKind.MOUNT)
|
|
|
|
def test_state_handler_helpers(context):
|
|
handler = MountHandler()
|
|
assert handler.definition().kind == WorkflowStateKind.MOUNT
|
|
|
|
context.current_state = None
|
|
assert handler.can_run(context) is True
|
|
context.current_state = WorkflowStateKind.MOUNT
|
|
assert handler.can_run(context) is True
|
|
context.current_state = WorkflowStateKind.RASTER
|
|
assert handler.can_run(context) is False
|
|
|
|
def test_workflow_runner_error_cases():
|
|
runner = WorkflowRunner()
|
|
with pytest.raises(KeyError, match="No handler registered"):
|
|
runner.get_handler("non_existent_state")
|
|
|
|
@patch("time.sleep", return_value=None)
|
|
def test_all_simulated_handlers(mock_sleep, context):
|
|
for kind, handler in SIMULATED_HANDLER_REGISTRY.items():
|
|
context.abort_requested = False
|
|
res = handler.execute(context)
|
|
assert res.status == StepStatus.SUCCESS
|
|
assert "SIMULATION" in res.message
|
|
|
|
context.abort_requested = True
|
|
with pytest.raises(RuntimeError, match="Abort requested"):
|
|
handler.execute(context)
|