140 lines
4.8 KiB
Python
140 lines
4.8 KiB
Python
import pytest
|
|
import asyncio
|
|
from unittest.mock import MagicMock, patch
|
|
from aare.daq.automation_runner import PersistentWorkflowRunner, AutomationLoop
|
|
from aare.common.automation_models import (
|
|
QueueItem, QueueItemStatus, WorkflowContext, WorkflowMode,
|
|
WorkflowStateKind, StateResult, StepStatus, ControlState, RuntimeState,
|
|
StepState
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_redis():
|
|
m = MagicMock()
|
|
m._bl = "x10sa"
|
|
m.get_runtime.return_value = RuntimeState()
|
|
m.get_control.return_value = ControlState()
|
|
return m
|
|
|
|
@pytest.fixture
|
|
def runner(mock_redis):
|
|
return PersistentWorkflowRunner(redis_manager=mock_redis)
|
|
|
|
def test_runner_init(runner, mock_redis):
|
|
assert runner.beamline == "x10sa"
|
|
assert runner._redis == mock_redis
|
|
|
|
def test_start_item(runner, mock_redis):
|
|
item = QueueItem(
|
|
item_id="item1", beamline="x10sa", sample_id=1, sample_name="S1",
|
|
owner_pgroup="p1", created_by="u", steps=[]
|
|
)
|
|
mock_redis.get_item.return_value = item
|
|
mock_redis.update_item.return_value = item
|
|
|
|
started_item = runner.start_item("item1")
|
|
|
|
assert started_item == item
|
|
mock_redis.update_item.assert_called_once()
|
|
mock_redis.patch_runtime.assert_called_once()
|
|
mock_redis.append_event.assert_called_once()
|
|
|
|
def test_run_current_step_success(runner, mock_redis):
|
|
step = MagicMock()
|
|
step.kind = "mount"
|
|
item = MagicMock(item_id="item1", steps=[step], current_step_index=0)
|
|
mock_redis.get_item.return_value = item
|
|
|
|
handler = MagicMock()
|
|
handler.execute.return_value = StateResult(
|
|
state=WorkflowStateKind.MOUNT, status=StepStatus.SUCCESS, message="OK"
|
|
)
|
|
runner._handlers = {WorkflowStateKind.MOUNT: handler}
|
|
|
|
context = WorkflowContext(
|
|
mode=WorkflowMode.AUTOMATION, queue_id="q", item_id="item1", sample_id=1
|
|
)
|
|
|
|
res = runner.run_current_step(context)
|
|
|
|
assert res.status == StepStatus.SUCCESS
|
|
mock_redis.update_step.assert_called()
|
|
mock_redis.update_item.assert_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_automation_loop_start_stop(runner, mock_redis, mocker):
|
|
loop = AutomationLoop(runner, mock_redis)
|
|
|
|
# Patch _run_loop to be a non-coroutine to avoid "never awaited" warning
|
|
mocker.patch.object(loop, "_run_loop", return_value=None)
|
|
|
|
with patch("asyncio.create_task") as mock_task:
|
|
mock_task.return_value = MagicMock()
|
|
loop.start()
|
|
assert loop.is_enabled is True
|
|
mock_task.assert_called_once()
|
|
|
|
loop.stop()
|
|
assert loop.is_enabled is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_automation_loop_tick_no_items(runner, mock_redis):
|
|
loop = AutomationLoop(runner, mock_redis)
|
|
mock_redis.get_runtime.return_value = RuntimeState(running=False)
|
|
mock_redis.get_next_pending_item.return_value = None
|
|
runner.check_control = MagicMock(return_value=ControlState())
|
|
|
|
# Ensure runner.start_item is a regular Mock, not AsyncMock
|
|
runner.start_item = MagicMock()
|
|
|
|
await loop._tick()
|
|
|
|
mock_redis.get_next_pending_item.assert_called_once()
|
|
mock_redis.patch_runtime.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_automation_loop_tick_process_step(runner, mock_redis):
|
|
loop = AutomationLoop(runner, mock_redis)
|
|
# 1. Mock runtime and control so it thinks an item is running
|
|
mock_redis.get_runtime.return_value = RuntimeState(
|
|
running=True, current_item_id="item1", current_state="mount"
|
|
)
|
|
runner.check_control = MagicMock(return_value=ControlState())
|
|
|
|
# 2. Mock item and its steps
|
|
step = MagicMock()
|
|
step.kind = "loop_centre"
|
|
item = MagicMock(item_id="item1", steps=[MagicMock(kind="mount"), step], current_step_index=1, sample_id=123)
|
|
mock_redis.get_item.return_value = item
|
|
|
|
# Ensure runner.run_current_step is a regular Mock
|
|
runner.run_current_step = MagicMock(return_value=StateResult(
|
|
state=WorkflowStateKind.LOOP_CENTRE, status=StepStatus.SUCCESS
|
|
))
|
|
|
|
await loop._tick()
|
|
|
|
runner.run_current_step.assert_called_once()
|
|
assert runner.run_current_step.call_args[0][0].item_id == "item1"
|
|
|
|
def test_runner_error_handling(runner, mock_redis):
|
|
step = MagicMock()
|
|
step.kind = "mount"
|
|
item = MagicMock(item_id="item1", steps=[step], current_step_index=0)
|
|
mock_redis.get_item.return_value = item
|
|
|
|
handler = MagicMock()
|
|
handler.execute.side_effect = Exception("Hardware failure")
|
|
runner._handlers = {WorkflowStateKind.MOUNT: handler}
|
|
|
|
context = WorkflowContext(
|
|
mode=WorkflowMode.AUTOMATION, queue_id="q", item_id="item1", sample_id=1
|
|
)
|
|
|
|
with pytest.raises(Exception, match="Hardware failure"):
|
|
runner.run_current_step(context)
|
|
|
|
mock_redis.update_step.assert_called_with(
|
|
"item1", 0, status="failed", error_detail="Hardware failure"
|
|
)
|