121 lines
3.2 KiB
Python
121 lines
3.2 KiB
Python
from dataclasses import asdict
|
|
|
|
from aare.common.automation_models import (
|
|
AutomationProgress,
|
|
StepState,
|
|
StepStatus,
|
|
WorkflowStateKind,
|
|
)
|
|
from aare.daq.config import BeamlineConfig
|
|
|
|
|
|
class _FakeRedis:
|
|
def __init__(self):
|
|
self._store: dict[str, str | int] = {}
|
|
|
|
def get(self, key: str):
|
|
return self._store.get(key)
|
|
|
|
def set(self, key: str, value):
|
|
self._store[key] = value
|
|
|
|
def delete(self, key: str):
|
|
self._store.pop(key, None)
|
|
|
|
def incr(self, key: str) -> int:
|
|
value = int(self._store.get(key, 0)) + 1
|
|
self._store[key] = value
|
|
return value
|
|
|
|
|
|
def _make_config_with_fake_redis() -> BeamlineConfig:
|
|
cfg = BeamlineConfig.__new__(BeamlineConfig)
|
|
cfg._BeamlineConfig__bl = "testbeamline"
|
|
cfg._BeamlineConfig__client = _FakeRedis()
|
|
return cfg
|
|
|
|
|
|
def test_automation_progress_state_round_trip_dataclass():
|
|
cfg = _make_config_with_fake_redis()
|
|
|
|
progress = AutomationProgress(
|
|
current_step="Center",
|
|
steps=[
|
|
StepState(
|
|
step=WorkflowStateKind.MOUNT,
|
|
status=StepStatus.SUCCESS,
|
|
message="Mount complete",
|
|
),
|
|
StepState(
|
|
step=WorkflowStateKind.LOOP_CENTRE,
|
|
status=StepStatus.RUNNING,
|
|
message="Centering sample",
|
|
),
|
|
StepState(
|
|
step=WorkflowStateKind.RASTER,
|
|
status=StepStatus.PENDING,
|
|
message="",
|
|
),
|
|
StepState(
|
|
step=WorkflowStateKind.DATA_COLLECTION,
|
|
status=StepStatus.PENDING,
|
|
message="",
|
|
),
|
|
StepState(
|
|
step=WorkflowStateKind.FINAL,
|
|
status=StepStatus.PENDING,
|
|
message="",
|
|
),
|
|
],
|
|
finished=False,
|
|
success=None,
|
|
)
|
|
|
|
written_state = cfg.set_automation_progress_state(progress)
|
|
|
|
assert written_state["seq"] == 1
|
|
assert written_state["progress"] == asdict(progress)
|
|
|
|
read_state = cfg.get_automation_progress_state()
|
|
|
|
assert read_state["seq"] == 1
|
|
assert read_state["progress"] == asdict(progress)
|
|
|
|
|
|
def test_automation_progress_state_seq_increments():
|
|
cfg = _make_config_with_fake_redis()
|
|
|
|
first = AutomationProgress(
|
|
current_step="Mount",
|
|
steps=[
|
|
StepState(
|
|
step=WorkflowStateKind.MOUNT,
|
|
status=StepStatus.RUNNING,
|
|
message="Mounting sample",
|
|
)
|
|
],
|
|
finished=False,
|
|
success=None,
|
|
)
|
|
second = AutomationProgress(
|
|
current_step="Paused/Finished",
|
|
steps=[
|
|
StepState(
|
|
step=WorkflowStateKind.FINAL,
|
|
status=StepStatus.SUCCESS,
|
|
message="Automation complete",
|
|
)
|
|
],
|
|
finished=True,
|
|
success=True,
|
|
)
|
|
|
|
first_state = cfg.set_automation_progress_state(first)
|
|
second_state = cfg.set_automation_progress_state(second)
|
|
|
|
assert first_state["seq"] == 1
|
|
assert second_state["seq"] == 2
|
|
|
|
final_state = cfg.get_automation_progress_state()
|
|
assert final_state["seq"] == 2
|
|
assert final_state["progress"] == asdict(second) |