127 lines
3.5 KiB
Python
127 lines
3.5 KiB
Python
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()
|