AareDAQ/tests/unit/daq/test_spreadsheetupdater.py

import json
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from aare.daq.spreadsheetupdater import get_ws_headers, on_message, set_spreadsheet_in_redis


@pytest.fixture
def mock_config():
    with patch("aare.daq.spreadsheetupdater.config") as mock:
        mock._bl = "X10SA"
        mock.redis = MagicMock()
        yield mock


def test_get_ws_headers_success():
    with patch("os.getenv", return_value="secret"):
        headers = get_ws_headers()
        assert headers == ["X-Shared-Password: secret"]


def test_get_ws_headers_fail():
    with patch("os.getenv", return_value=None), pytest.raises(ValueError):
        get_ws_headers()


def test_set_spreadsheet_in_redis(mock_config):
    data = {"test": "data"}
    with patch("aare.daq.spreadsheetupdater.config") as mock_cfg_internal:
        mock_client = MagicMock()
        mock_cfg_internal.redis = mock_client

        set_spreadsheet_in_redis(data)

        found = False
        for attr in dir(mock_cfg_internal):
            val = getattr(mock_cfg_internal, attr)
            if isinstance(val, MagicMock) and val.set.called:
                found = True
                break
        assert found or mock_client.set.called


def test_on_message_success(mock_config):
    normal_sample = SimpleNamespace(
        id=1,
        sample_name="S1",
        run_number=100,
        pgroup="p12345",
        position=1,
        priority=1,
        mount_count=0,
        rotation_count=0,
        raster_count=0,
        screening_count=0,
        data_collection_parameters={},
    )
    ref_sample = SimpleNamespace(
        id=2,
        sample_name="R1",
        run_number=1,
        pgroup="p12345",
        position=1,
        priority=1,
        mount_count=0,
        rotation_count=0,
        raster_count=0,
        screening_count=0,
        data_collection_parameters={},
    )

    mock_pucks = [
        SimpleNamespace(
            puck_name="P1", dewar_name="D1", tell_position="A1", samples=[normal_sample]
        ),
        SimpleNamespace(puck_name="Ref", dewar_name="D1", tell_position="X1", samples=[ref_sample]),
    ]

    message = json.dumps({"samples": [{}, {}]})

    with patch("aare.daq.spreadsheetupdater.PuckWithTellPosition", side_effect=mock_pucks):
        on_message(None, message)

    calls = mock_config.redis.set.call_args_list
    written_keys = [call.args[0] for call in calls]

    assert "X10SA:sample_spreadsheet" in written_keys
    assert "X10SA:reference-tools" in written_keys


def test_on_message_empty_ref(mock_config):
    message = json.dumps(
        {
            "samples": [
                {
                    "id": 1,
                    "barcode": "B1",
                    "position": "P1",
                    "puck_name": "P1",
                    "puck_type": "UniPuck",
                    "puck_location_in_dewar": 1,
                    "dewar_id": 1,
                    "pgroup": "p12345",
                    "dewar_name": "D1",
                    "tell_position": "A1",
                    "samples": [],
                }
            ]
        }
    )

    on_message(None, message)

    ref_key = "X10SA:reference-tools"
    mock_config.redis.delete.assert_called_with(ref_key)


def test_on_message_invalid_json(mock_config):
    on_message(None, "invalid json")
    mock_config.redis.set.assert_not_called()