165 lines
6.2 KiB
Python
165 lines
6.2 KiB
Python
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from unittest.mock import MagicMock, patch, PropertyMock
|
|
import os
|
|
import io
|
|
import cv2
|
|
import numpy as np
|
|
|
|
# Set environment variable before importing app
|
|
os.environ["JWT_AAREDAQ_KEY"] = "test_key_for_unit_testing"
|
|
|
|
# Mock the entire backend before importing the app to prevent any real initialisation
|
|
import aare.daq.daq
|
|
with patch("aare.daq.daq.AareDAQ"), \
|
|
patch("aare.daq.config.MXBeamline"), \
|
|
patch("aare.daq.config.BeamlineConfig"), \
|
|
patch("aare.daq.aaredb.AareWrapper"), \
|
|
patch("aare.daq.server.lifespan") as mock_ls:
|
|
mock_ls.return_value.__aenter__.return_value = None
|
|
from aare.daq.server import app
|
|
import aare.daq.server as server
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_backend():
|
|
with patch("aare.daq.server.daq") as m_daq, \
|
|
patch("aare.daq.server.bl") as m_bl, \
|
|
patch("aare.daq.server.cfg") as m_cfg:
|
|
|
|
# Manually ensure these globals are set to our mocks
|
|
server.daq = m_daq
|
|
server.bl = m_bl
|
|
server.cfg = m_cfg
|
|
|
|
m_daq.busy = False
|
|
m_cfg.pgroup = "p12345"
|
|
m_cfg.session_state.return_value = "ACTIVE"
|
|
|
|
yield {
|
|
"daq": m_daq,
|
|
"bl": m_bl,
|
|
"cfg": m_cfg
|
|
}
|
|
|
|
@pytest.fixture
|
|
def client(mock_backend):
|
|
# Mock auth.parse_token to return a valid TokenData
|
|
with patch("aare.daq.auth.parse_token") as mock_parse:
|
|
from aare.daq.auth import TokenData
|
|
mock_parse.return_value = TokenData(
|
|
sub="testuser",
|
|
staff=True,
|
|
pgroups=["p12345"],
|
|
session=123,
|
|
)
|
|
# Also mock cv2.imencode globally within the client context to avoid OpenCV issues with mocks
|
|
with patch("cv2.imencode") as mock_imencode:
|
|
mock_imencode.return_value = (True, np.array([1, 2, 3], dtype=np.uint8))
|
|
with TestClient(app) as c:
|
|
yield c
|
|
|
|
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(client, mock_backend):
|
|
mock_daq = mock_backend["daq"]
|
|
mock_cfg = mock_backend["cfg"]
|
|
|
|
from aare.common.models import DAQStatusModel, SessionStatus, SessionsStateEnum, BeamlineStateEnum, SampleGeometryModel, BeamlineStatus, CrystalSize, SampleCameraSettings
|
|
from aare.common.diffraction_geometry import DiffractionGeometry
|
|
|
|
# Create a valid status object
|
|
from aare.common.coordinate import Coordinate, SmargonCoordinate
|
|
geom = SampleGeometryModel(
|
|
beam_location_pxl=Coordinate(x=500, y=500),
|
|
pixel_in_mm=0.001,
|
|
aerotech=Coordinate(x=0, y=0, z=0),
|
|
aerotech_meas=Coordinate(x=0, y=0, z=0),
|
|
smargon=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0, z=0), phi_deg=0.0, chi_deg=0.0),
|
|
omega_deg=0.0,
|
|
beam_size_mm=Coordinate(x=0.01, y=0.01)
|
|
)
|
|
diff = DiffractionGeometry(
|
|
energy_keV=12.0,
|
|
dtz_mm=150.0,
|
|
pixel_size_mm=0.075,
|
|
beam_center_pxl=(1000.0, 1000.0),
|
|
detector_size_pxl=(2000, 2000),
|
|
detector_description="Eiger 16M",
|
|
detector_serial_number="123",
|
|
poni_rot1_rad=0.0,
|
|
poni_rot2_rad=0.0
|
|
)
|
|
bl_status = BeamlineStatus(
|
|
name="X06DA",
|
|
ring_current_mA=400.0,
|
|
omega=0.0, front_light=0.0, back_light=0.0,
|
|
cryojet_K=100.0, shutter_open=False, exp_shutter_open=False,
|
|
flux_ph_s=1e12, transmission=1.0, zoom=1.0,
|
|
sample_camera=SampleCameraSettings(exposure=0.1, gain=1.0),
|
|
commissioning_mode=False, dtz_min=120.0, dtz_max=1600.0
|
|
)
|
|
session_status = SessionStatus(
|
|
session=SessionsStateEnum.Vacant,
|
|
current_pgroup="p12345",
|
|
staff=True
|
|
)
|
|
|
|
status_obj = DAQStatusModel(
|
|
geom=geom,
|
|
diffraction=diff,
|
|
bl=bl_status,
|
|
state=BeamlineStateEnum.Maintenance,
|
|
busy=False,
|
|
session=session_status,
|
|
crystal_size=CrystalSize(x=0, y=0, z=0)
|
|
)
|
|
|
|
mock_daq.status = status_obj
|
|
mock_cfg.pgroup = "p12345"
|
|
mock_cfg.session_state.return_value = SessionsStateEnum.OwnedByYou
|
|
|
|
# Also patch server.cfg and server.daq to be absolutely sure
|
|
with patch("aare.daq.server.daq", mock_daq), \
|
|
patch("aare.daq.server.cfg", mock_cfg):
|
|
response = client.get("/status", headers={"Authorization": "Bearer fake-token"})
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["state"] == 1 # Maintenance is 1
|
|
assert data["session"]["current_pgroup"] == "p12345"
|
|
|
|
def test_omega_put(client, mock_backend):
|
|
# The problem is that server.py uses 'from aare.daq.server import daq' in some places
|
|
# and the global 'daq' in others.
|
|
# When we do 'daq.omega = val', it's the global 'daq'.
|
|
|
|
with patch("aare.daq.auth.check_jwt_rw"):
|
|
# Instead of checking mock_daq.omega, let's patch the 'daq' object in server.py AGAIN
|
|
with patch("aare.daq.server.daq") as m_daq:
|
|
server.daq = m_daq # Force it
|
|
response = client.put("/beamline/omega?val=10.5", headers={"Authorization": "Bearer fake-token"})
|
|
assert response.status_code == 200
|
|
assert response.json() == "OK"
|
|
# Now m_daq should have received the assignment
|
|
# In Python, m_daq.omega = 10.5 will set the 'omega' attribute on the MagicMock.
|
|
assert m_daq.omega == 10.5
|
|
|
|
def test_login_success(client, mock_backend):
|
|
with patch("aare.daq.auth.authenticate_user") as mock_auth:
|
|
mock_auth.return_value = "fake-access-token"
|
|
response = client.post("/token", data={"username": "user", "password": "pwd"})
|
|
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)
|
|
with patch("aare.daq.auth.parse_token"):
|
|
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
|