update test cases for workflows, server and daq_server

tests/integration/daq/test_daq_server.py

@@ -25,6 +25,22 @@ def test_read_error_codes(client):
 @pytest.mark.integration
 def test_login_unauthorized(client):
     # Testing login with invalid credentials. 
-    # The current implementation raises KeyError if user not found.
-    with pytest.raises(Exception):
-        client.post("/token", data={"username": "non_existent_user_123", "password": "bad"})
+    # The current implementation raises KeyError if user not found, 
+    # which FastAPI might convert to 500 or just propagate if using TestClient in some modes.
+    # However, let's just check for a non-200 status code.
+    response = client.post("/token", data={"username": "non_existent_user_123", "password": "bad"})
+    assert response.status_code != 200
+
+
+@pytest.mark.integration
+def test_status_unauthorized(client):
+    # Should fail because no Bearer token is provided
+    response = client.get("/status")
+    assert response.status_code == 401
+
+
+@pytest.mark.integration
+def test_pgroup_unauthorized(client):
+    # Should fail because no Bearer token is provided
+    response = client.get("/access/pgroup")
+    assert response.status_code == 401

tests/unit/daq/test_server.py

@@ -1,12 +1,11 @@
 import os
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 
 os.environ["JWT_AAREDAQ_KEY"] = "test_key_for_unit_testing"
 
-
 def test_meta_error_codes(client):
     response = client.get("/meta/error-codes")
     assert response.status_code == 200
@@ -14,14 +13,24 @@ def test_meta_error_codes(client):
     assert "AuthErrorCode" in data
 
 
-def test_status(api, daq_status_factory):
+def test_status(api, daq_status_factory, monkeypatch):
     from aare.common.models import BeamlineStateEnum, SessionsStateEnum
+    from aare.daq import server
+
+    monkeypatch.setattr(server.auth, "resolve_baton_timeout_if_needed", lambda cfg: None)
+    monkeypatch.setattr(server.auth, "get_baton_status", lambda cfg, data: {"dummy": "status"})
+
+    api.cfg.pending_baton_request = None
+    api.cfg.queued_baton_transfer = None
+    api.cfg.baton_holder = None
+    api.cfg.allow_non_staff_request_from_staff = False
 
     api.daq.status = daq_status_factory(
         state=BeamlineStateEnum.Maintenance,
         current_pgroup="p12345",
         staff=True,
     )
+    # ... rest of the setup
     api.daq.status.sample = None
     api.daq.status.box = None
     api.daq.status.last_best_res = None
@@ -67,4 +76,202 @@ def test_get_image(client, mock_backend):
     response = client.get("/beamline/image", headers={"Authorization": "Bearer fake-token"})
     assert response.status_code == 200
     assert response.headers["content-type"] == "image/jpeg"
-    assert len(response.content) > 0
+    assert len(response.content) > 0
+
+
+def test_mount_returns_tell_exception_when_mount_precheck_fails(client, mock_backend, monkeypatch):
+    from aare.common.exception_handler import TellCommunicationError
+    from aare.daq import server
+
+    monkeypatch.setattr(server.auth, "check_jwt_rw", lambda *_args, **_kwargs: None)
+
+    mock_daq = mock_backend["daq"]
+    mock_daq.check_tell_mount_start_conditions.side_effect = TellCommunicationError(
+        "Mount can't start: system check failed: Exception('Psys safety not released',); TELL doors are open"
+    )
+
+    response = client.post(
+        "/sample/mount?dbid=1&reference=false",
+        headers={"Authorization": "Bearer fake-token"},
+    )
+
+    assert response.status_code == 503
+    payload = response.json()
+    assert payload["exception_class"] == "TellCommunicationError"
+    assert payload["code"] == "TELL_COMMUNICATION_ERROR"
+    assert "Mount can't start:" in payload["message"]
+
+
+def test_mount_calls_tell_mount_precheck_before_mount(client, mock_backend, monkeypatch):
+    from aare.daq import server
+
+    monkeypatch.setattr(server.auth, "check_jwt_rw", lambda *_args, **_kwargs: None)
+
+    mock_daq = mock_backend["daq"]
+
+    sample = SimpleNamespace(db_id=1, user="p12345")
+    mock_daq.sample_spreadsheet = SimpleNamespace(s=[sample])
+    mock_daq.reference_tools = SimpleNamespace(s=[])
+
+    response = client.post(
+        "/sample/mount?dbid=1&reference=false",
+        headers={"Authorization": "Bearer fake-token"},
+    )
+
+    assert response.status_code == 200
+    mock_daq.check_tell_mount_start_conditions.assert_called_once_with()
+    assert mock_daq.sample == sample
+
+
+def test_auto_scan_returns_tell_exception_when_mount_precheck_fails(client, mock_backend, monkeypatch):
+    from aare.common.exception_handler import TellCommunicationError
+    from aare.daq import server
+
+    monkeypatch.setattr(server.auth, "check_jwt_rw", lambda *_args, **_kwargs: None)
+
+    mock_daq = mock_backend["daq"]
+    mock_daq.check_tell_mount_start_conditions.side_effect = TellCommunicationError(
+        "Mount can't start: TELL is not in remote mode; TELL doors are open"
+    )
+
+    body = {
+        "db_id": 1,
+        "puck_name": "puck1",
+        "dewar_name": "dewar1",
+        "sample_name": "sample1",
+        "run_number": 1,
+        "user": "p12345",
+        "pin": 1,
+        "location": {"segment": "A", "pos": 1},
+    }
+
+    response = client.post(
+        "/scan/auto",
+        json=body,
+        headers={"Authorization": "Bearer fake-token"},
+    )
+
+    assert response.status_code == 503
+    payload = response.json()
+    assert payload["exception_class"] == "TellCommunicationError"
+    assert payload["code"] == "TELL_COMMUNICATION_ERROR"
+    assert "Mount can't start:" in payload["message"]
+
+
+def test_auto_scan_calls_tell_mount_precheck_before_measure(client, mock_backend, monkeypatch):
+    from aare.daq import server
+
+    monkeypatch.setattr(server.auth, "check_jwt_rw", lambda *_args, **_kwargs: None)
+
+    mock_daq = mock_backend["daq"]
+    mock_daq.measure.return_value = 12.345
+
+    body = {
+        "db_id": 1,
+        "puck_name": "puck1",
+        "dewar_name": "dewar1",
+        "sample_name": "sample1",
+        "run_number": 1,
+        "user": "p12345",
+        "pin": 1,
+        "location": {"segment": "A", "pos": 1},
+    }
+
+    response = client.post(
+        "/scan/auto",
+        json=body,
+        headers={"Authorization": "Bearer fake-token"},
+    )
+
+    assert response.status_code == 200
+    mock_daq.check_tell_mount_start_conditions.assert_called_once_with()
+    mock_daq.measure.assert_called_once()
+
+
+def test_get_pgroup(api):
+    api.cfg.pgroup = "p12345"
+    response = api.client.get("/access/pgroup", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == "p12345"
+
+
+def test_set_pgroup(api):
+    api.cfg.baton_holder = None
+    response = api.client.put("/access/pgroup?val=p54321", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == "OK"
+    assert api.cfg.pgroup == "p54321"
+
+
+def test_delete_pgroup(api):
+    api.client.delete("/access/pgroup", headers={"Authorization": "Bearer fake-token"})
+    assert api.cfg.pgroup is None
+
+
+def test_set_commissioning_mode(api):
+    response = api.client.put("/beamline/commissioning_mode?val=true", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == "OK"
+    assert api.cfg.commissioning_mode is True
+
+
+def test_get_settings(api):
+    from aare.common.models import BeamlineSettingsModel
+    mock_settings = BeamlineSettingsModel()
+    api.cfg.settings = mock_settings
+    response = api.client.get("/beamline/settings", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == mock_settings.model_dump()
+
+
+def test_put_settings(api):
+    from aare.common.models import BeamlineSettingsModel
+    settings_data = BeamlineSettingsModel().model_dump()
+    response = api.client.put("/beamline/settings", json=settings_data, headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert api.cfg.settings.model_dump() == settings_data
+
+
+def test_get_cryo_settings(api):
+    from aare.common.models import CryojetSettingsModel
+    mock_cryo = CryojetSettingsModel()
+    api.cfg.cryojet_settings = mock_cryo
+    response = api.client.get("/beamline/cryo_settings", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == mock_cryo.model_dump()
+
+
+def test_put_cryo_settings(api):
+    from aare.common.models import CryojetSettingsModel
+    cryo_data = CryojetSettingsModel().model_dump()
+    response = api.client.put("/beamline/cryo_settings", json=cryo_data, headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert api.cfg.cryojet_settings.model_dump() == cryo_data
+
+
+def test_baton_status(api):
+    from aare.common.auth_models import BatonStatus
+    mock_baton = BatonStatus(holder=None, request=None, allow_non_staff_request=True)
+    api.cfg.baton_status = mock_baton
+    api.cfg.baton_holder = None
+    api.cfg.queued_baton_transfer = None
+    api.cfg.allow_non_staff_request_from_staff = True
+    response = api.client.get("/baton/status", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == mock_baton.model_dump()
+
+
+def test_baton_request(api, monkeypatch):
+    from aare.daq import server
+    monkeypatch.setattr(server.auth, "request_baton", lambda cfg, data: {"granted": True})
+    response = api.client.post("/baton/request", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == {"granted": True}
+
+
+def test_baton_release(api, monkeypatch):
+    from aare.daq import server
+    monkeypatch.setattr(server.auth, "release_baton", lambda cfg, data: {"released": True})
+    response = api.client.post("/baton/release", headers={"Authorization": "Bearer fake-token"})
+    assert response.status_code == 200
+    assert response.json() == {"released": True}

tests/unit/daq/test_workflows.py

@@ -1,6 +1,8 @@
+from types import SimpleNamespace
+
 import pytest
 from unittest.mock import MagicMock, patch
-from aare.daq.workflows import common_2rse, sa2se, sa2rse, sa2xtal_snapshot, dc2xtal_snapshot, xtal_snapshot2dc, xtal_snapshot2sa, dc2rse, se2sa, sa2dc, dc2sa, sa2xrf, sa2dh, dh2sa
+from aare.daq.workflows import common_2rse, sa2se, sa2rse, sa2xtal_snapshot, dc2xtal_snapshot, xtal_snapshot2dc, xtal_snapshot2sa, dc2rse, se2sa, sa2dc, dc2sa, sa2xrf, sa2dh, dh2sa, common2dh
 from aare.daq.config import ABR_POS_MOUNT, ABR_OMEGA_MOUNT
 from aare.common.models import StagePositionEnum
 from aare.devices.area_detector import AutoEnum
@@ -205,4 +207,20 @@ def test_dh2sa(mock_devs, mock_cfg):
     mock_devs.bec_worker.move_to = MagicMock()
 
     dh2sa(mock_devs, mock_cfg)
-    _assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
+    _assert_bec_moved(mock_devs, BeamlineState.SAMPLE_ALIGNMENT)
+
+
+def test_common2dh_skips_dry_when_tell_already_in_ppark(mock_devs, mock_cfg):
+    mock_devs.bec_worker = MagicMock()
+    mock_devs.bec_worker.planner = MagicMock()
+    mock_devs.bec_worker.move_to = MagicMock()
+    mock_devs.tell = MagicMock()
+    mock_devs.tell.is_position.return_value = True
+    mock_devs.tell.get_mounted_sample.return_value = None
+
+    common2dh(mock_devs, mock_cfg)
+
+    mock_devs.tell.is_position.assert_called_once_with("pPark")
+    mock_devs.tell.get_mounted_sample.assert_called_once_with()
+    mock_devs.tell.unmount.assert_not_called()
+    mock_devs.tell.dry.assert_not_called()