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._client = MagicMock() # Mocking private attributes access which the code uses mock._client = mock._client mock._bl = mock._bl 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._client = mock_client mock_cfg_internal.client = 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._client.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._client.delete.assert_called_with(ref_key) def test_on_message_invalid_json(mock_config): on_message(None, "invalid json") mock_config._client.set.assert_not_called()