diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index 9bd31011..074e1897 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -1003,6 +1003,63 @@ async def beam_location(token: str = Depends(oauth2_scheme)): auth.check_jwt_staff(cfg, auth.parse_token(token)) daq.state = BeamlineStateEnum.BeamLocation +@app.post("/state/data_collection") +async def data_collection(token: str = Depends(oauth2_scheme)): + """ + Transition beamline state to DewarTransfer. + + Args: + token: OAuth2 access token. + + Returns: + "OK" on success. + """ + auth.check_jwt_rw(cfg, auth.parse_token(token)) + daq.state = BeamlineStateEnum.DataCollection + return "OK" + + +@app.post("/state/robot_sample_exchange") +async def dewar_exchange(token: str = Depends(oauth2_scheme)): + """ + Transition beamline state to SampleExchange. + + Args: + token: OAuth2 access token. + + Returns: + "OK" on success. + """ + auth.check_jwt_rw(cfg, auth.parse_token(token)) + daq.state = BeamlineStateEnum.RobotSampleExchange + return "OK" + + +@app.post("/state/xray_fluorescence") +async def xray_fluorescence(token: str = Depends(oauth2_scheme)): + """ + Transition beamline state to SampleAlignment. + + Args: + token: OAuth2 access token. + + Returns: + "OK" on success. + """ + auth.check_jwt_rw(cfg, auth.parse_token(token)) + daq.state = BeamlineStateEnum.XrayFluorescence + +@app.post("/state/xtal_snapshot") +async def beam_location(token: str = Depends(oauth2_scheme)): + """ + Transition beamline state to BeamLocation. Staff only. + + Args: + token: OAuth2 access token. + """ + auth.check_jwt_staff(cfg, auth.parse_token(token)) + daq.state = BeamlineStateEnum.XtalSnapshot + @app.post("/state/maintenance") async def maintenance(token: str = Depends(oauth2_scheme)) -> str: """ diff --git a/src/aare/gui/panels/beamline_state_panel.py b/src/aare/gui/panels/beamline_state_panel.py index 7da5e1ef..4f036609 100644 --- a/src/aare/gui/panels/beamline_state_panel.py +++ b/src/aare/gui/panels/beamline_state_panel.py @@ -52,6 +52,7 @@ class BeamlineStatePanel(QFrame): data_collection = Signal() xtal_snapshot = Signal() xray_fluorescence = Signal() + robot_sample_exchange = Signal() set_width = 400 map_height = 560 @@ -111,7 +112,7 @@ class BeamlineStatePanel(QFrame): self._metro_stations = [ StationSpec(BeamlineStateEnum.DewarTransfer, "Dewar transfer", 96, 126, True, "Dewar transfer mode"), StationSpec(BeamlineStateEnum.SampleExchange, "Manual sample exchange", 96, 186, True, "Manual sample exchange mode"), - StationSpec(BeamlineStateEnum.RobotSampleExchange, "Robot sample exchange", 96, 246, False, "Robot-assisted sample exchange"), + StationSpec(BeamlineStateEnum.RobotSampleExchange, "Robot sample exchange", 96, 246, True, "Robot-assisted sample exchange"), StationSpec(BeamlineStateEnum.SampleAlignment, "Sample alignment", 96, 316, True, "Sample centring and alignment mode"), StationSpec(BeamlineStateEnum.BeamLocation, "Beam location", 250, 316, True, "Beam location mode"), StationSpec(BeamlineStateEnum.DataCollection, "Data collection", 96, 396, True, "Measurement / collection mode"), @@ -122,7 +123,7 @@ class BeamlineStatePanel(QFrame): self._line_stations = [ StationSpec(BeamlineStateEnum.DewarTransfer, "Dewar transfer", 54, 170, True, "Dewar transfer mode"), StationSpec(BeamlineStateEnum.SampleExchange, "Manual sample exchange", 54, 220, True, "Manual sample exchange mode"), - StationSpec(BeamlineStateEnum.RobotSampleExchange, "Robot sample exchange", 54, 270, False, "Robot-assisted sample exchange"), + StationSpec(BeamlineStateEnum.RobotSampleExchange, "Robot sample exchange", 54, 270, True, "Robot-assisted sample exchange"), StationSpec(BeamlineStateEnum.SampleAlignment, "Sample alignment", 54, 320, True, "Sample centring and alignment mode"), StationSpec(BeamlineStateEnum.BeamLocation, "Beam location", 54, 370, True, "Beam location mode"), StationSpec(BeamlineStateEnum.DataCollection, "Data collection", 54, 420, True, "Measurement / collection mode"), @@ -330,6 +331,8 @@ class BeamlineStatePanel(QFrame): self.sample_alignment.emit() elif state == BeamlineStateEnum.DewarTransfer: self.dewar_exchange.emit() + elif state == BeamlineStateEnum.RobotSampleExchange: + self.robot_sample_exchange.emit() elif state == BeamlineStateEnum.BeamLocation: self.beam_location.emit() elif state == BeamlineStateEnum.DataCollection: diff --git a/src/aare/gui/threads/daq_worker.py b/src/aare/gui/threads/daq_worker.py index e90cea7d..d2d12638 100644 --- a/src/aare/gui/threads/daq_worker.py +++ b/src/aare/gui/threads/daq_worker.py @@ -695,6 +695,22 @@ class DAQWorker(QObject): def beam_location(self): self.generic_post("state/beam_location") + @Slot() + def data_collection(self): + self.generic_post("state/data_collection") + + @Slot() + def robot_sample_exchange(self): + self.generic_post("state/robot_sample_exchange") + + @Slot() + def xray_fluorescence(self): + self.generic_post("state/xray_fluorescence") + + @Slot() + def xtal_snapshot(self): + self.generic_post("state/xtal_snapshot") + @Slot(str) def free_beamline(self, confirmation_code: str): if self.__base_url is None: diff --git a/tests/unit/daq/test_automation_progress_state_manager.py b/tests/unit/daq/test_automation_progress_state_manager.py new file mode 100644 index 00000000..46796e52 --- /dev/null +++ b/tests/unit/daq/test_automation_progress_state_manager.py @@ -0,0 +1,121 @@ +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) \ No newline at end of file diff --git a/tests/unit/gui/test_automation_progress_parser.py b/tests/unit/gui/test_automation_progress_parser.py new file mode 100644 index 00000000..9564e7a7 --- /dev/null +++ b/tests/unit/gui/test_automation_progress_parser.py @@ -0,0 +1,82 @@ +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"