294 lines
10 KiB
Python
294 lines
10 KiB
Python
import os
|
|
from types import SimpleNamespace
|
|
from unittest.mock import 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
|
|
data = response.json()
|
|
assert "AuthErrorCode" in data
|
|
|
|
|
|
def test_status(api, daq_status_factory, monkeypatch):
|
|
from aarecommon.models.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
|
|
api.daq.status.last_best_b_factor = None
|
|
api.daq.status.crystal_size = SimpleNamespace(x=0, y=0, z=0)
|
|
api.daq.status.open_guis = []
|
|
|
|
api.cfg.pgroup = "p12345"
|
|
api.cfg.session_state.return_value = SessionsStateEnum.OwnedByYou
|
|
api.cfg.get_open_gui_sessions.return_value = []
|
|
|
|
response = api.client.get("/status", headers={"Authorization": "Bearer fake-token"})
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["state"] == BeamlineStateEnum.Maintenance.value
|
|
assert data["session"]["current_pgroup"] == "p12345"
|
|
assert data["session"]["session"] == SessionsStateEnum.OwnedByYou.value
|
|
|
|
|
|
def test_omega_put(client):
|
|
with patch("aare.daq.auth.check_jwt_rw"), patch("aare.daq.server.daq") as mock_daq:
|
|
response = client.put(
|
|
"/beamline/omega?val=10.5", headers={"Authorization": "Bearer fake-token"}
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json() == "OK"
|
|
assert mock_daq.omega == 10.5
|
|
|
|
|
|
def test_login_success(client):
|
|
with (
|
|
patch("aare.daq.auth.authenticate_from_proxy_header", return_value="user"),
|
|
patch("aare.daq.auth.authenticate_user", return_value="fake-access-token"),
|
|
):
|
|
response = client.post(
|
|
"/token",
|
|
data={"username": "user", "password": "pwd"},
|
|
headers={"X-Remote-User": "user"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json() == {"access_token": "fake-access-token", "token_type": "bearer"}
|
|
|
|
|
|
def test_get_image(client, mock_backend):
|
|
mock_daq = mock_backend["daq"]
|
|
mock_daq.camera_image = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
|
|
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
|
|
|
|
|
|
def test_mount_returns_tell_exception_when_mount_precheck_fails(client, mock_backend, monkeypatch):
|
|
from aarecommon.errors.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 aarecommon.errors.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 aarecommon.models.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 aarecommon.models.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 aarecommon.models.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 aarecommon.models.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 aarecommon.models.auth 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}
|