AareDAQ/tests/conftest.py

import os
import sys
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
from PySide6.QtWidgets import QApplication
from fastapi.testclient import TestClient

from aare.common.models import SampleShortInfo, DewarAddress

# Ensure safe defaults during tests
os.environ.setdefault("BEAMLINE", "SIMULATED")
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
os.environ.setdefault("JWT_AAREDAQ_KEY", "test_key_for_unit_testing")


@pytest.fixture
def mock_logger():
    return MagicMock()


# -------------------------
# Qt app fixture
# -------------------------
@pytest.fixture(scope="session")
def qapp():
    app = QApplication.instance()
    if app is None:
        app = QApplication(sys.argv)
    return app


@pytest.fixture
def sample_info():
    return SampleShortInfo(
        db_id=1,
        puck_name="puck1",
        dewar_name="dew1",
        sample_name="sample1",
        run_number=1,
        user="group1",
        pin=3,
        location=DewarAddress(segment="A", pos=2),
    )


# -------------------------
# Shared DAQ server fixtures
# -------------------------
@pytest.fixture(scope="session")
def server_module():
    import aare.daq.server as server
    return server


@pytest.fixture
def mock_backend(server_module):
    with patch.object(server_module, "daq", create=True) as m_daq, \
         patch.object(server_module, "bl", create=True) as m_bl, \
         patch.object(server_module, "cfg", create=True) as m_cfg:

        m_daq.busy = False
        m_daq.camera_image = np.zeros((10, 10, 3), dtype=np.uint8)

        m_cfg.pgroup = "p12345"
        m_cfg.session_state.return_value = "ACTIVE"
        m_cfg.GUI_SESSION_EXPIRE_SECONDS = 10
        m_cfg.get_open_gui_sessions.return_value = []
        m_cfg.get_gui_session.return_value = None
        m_cfg.touch_gui_session.return_value = None

        yield {
            "daq": m_daq,
            "bl": m_bl,
            "cfg": m_cfg,
        }


@pytest.fixture
def auth_token_data():
    from aare.daq.auth import TokenData

    return TokenData(
        sub="testuser",
        staff=True,
        pgroups=["p12345"],
        session=123,
    )


@pytest.fixture
def client(server_module, mock_backend, auth_token_data):
    with patch("aare.daq.auth.parse_token", return_value=auth_token_data), \
         patch("cv2.imencode", return_value=(True, np.array([1, 2, 3], dtype=np.uint8))), \
         patch.object(server_module, "mx_beamline", return_value=mock_backend["bl"]), \
         patch.object(server_module, "BeamlineConfig", return_value=mock_backend["cfg"]), \
         patch.object(server_module, "AareDAQ", return_value=mock_backend["daq"]):
        with TestClient(server_module.app) as c:
            yield c


@pytest.fixture
def api(client, mock_backend):
    return SimpleNamespace(
        client=client,
        daq=mock_backend["daq"],
        bl=mock_backend["bl"],
        cfg=mock_backend["cfg"],
    )


@pytest.fixture
def daq_status_factory():
    from aare.common.coordinate import Coordinate, SmargonCoordinate
    from aare.common.diffraction_geometry import DiffractionGeometry
    from aare.common.models import (
        BeamlineStateEnum,
        BeamlineStatus,
        CrystalSize,
        DAQStatusModel,
        SampleCameraSettings,
        SampleGeometryModel,
        SessionStatus,
        SessionsStateEnum,
    )

    def _build(
        *,
        state=BeamlineStateEnum.Maintenance,
        busy=False,
        current_pgroup="p12345",
        staff=True,
        session_state=SessionsStateEnum.Vacant,
        sample=None,
        open_guis=None,
    ):
        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 = SessionStatus(
            session=session_state,
            current_pgroup=current_pgroup,
            staff=staff,
        )

        status = DAQStatusModel(
            geom=geom,
            diffraction=diff,
            bl=bl_status,
            state=state,
            busy=busy,
            session=session,
            crystal_size=CrystalSize(x=0, y=0, z=0),
            sample=sample,
        )

        if open_guis is not None:
            status.open_guis = open_guis

        return status

    return _build