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 aarecommon.models.models import DewarAddress, SampleShortInfo
from fastapi.testclient import TestClient
from PySide6.QtWidgets import QApplication

# 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 aarecommon.math.coordinate import Coordinate, SmargonCoordinate
    from aarecommon.math.diffraction_geometry import DiffractionGeometry
    from aarecommon.models.models import (
        BeamlineStateEnum,
        BeamlineStatus,
        CrystalSize,
        DAQStatusModel,
        SampleCameraSettings,
        SampleGeometryModel,
        SessionsStateEnum,
        SessionStatus,
    )

    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