140 lines
4.4 KiB
Python
140 lines
4.4 KiB
Python
import pytest
|
|
import json
|
|
from unittest.mock import MagicMock, patch
|
|
from aare.daq.spreadsheetupdater import on_message, get_ws_headers, set_spreadsheet_in_redis
|
|
from aare.common.models import SampleShortInfoList
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_config():
|
|
with patch('aare.daq.spreadsheetupdater.config') as mock:
|
|
mock._BeamlineConfig__bl = "X10SA"
|
|
mock._BeamlineConfig__client = MagicMock()
|
|
# Mocking private attributes access which the code uses
|
|
mock._BeamlineConfig__client = mock._BeamlineConfig__client
|
|
mock._BeamlineConfig__bl = mock._BeamlineConfig__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):
|
|
with 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._BeamlineConfig__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):
|
|
message = json.dumps({
|
|
"samples": [
|
|
{
|
|
"id": 1,
|
|
"puck_name": "P1",
|
|
"puck_type": "UniPuck",
|
|
"puck_location_in_dewar": 1,
|
|
"dewar_id": 1,
|
|
"pgroup": "p12345",
|
|
"dewar_name": "D1",
|
|
"tell_position": "A1",
|
|
"samples": [
|
|
{
|
|
"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": {}
|
|
}
|
|
]
|
|
},
|
|
{
|
|
"id": 2,
|
|
"puck_name": "Ref",
|
|
"puck_type": "UniPuck",
|
|
"puck_location_in_dewar": 2,
|
|
"dewar_id": 1,
|
|
"pgroup": "p12345",
|
|
"dewar_name": "D1",
|
|
"tell_position": "X1",
|
|
"samples": [
|
|
{
|
|
"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": {}
|
|
}
|
|
]
|
|
}
|
|
]
|
|
})
|
|
|
|
on_message(None, message)
|
|
|
|
normal_key = "X10SA:sample_spreadsheet"
|
|
calls = mock_config._BeamlineConfig__client.set.call_args_list
|
|
assert any(call.args[0] == normal_key for call in calls)
|
|
|
|
ref_key = "X10SA:reference-tools"
|
|
assert any(call.args[0] == ref_key for call in calls)
|
|
|
|
|
|
def test_on_message_empty_ref(mock_config):
|
|
message = json.dumps({
|
|
"samples": [
|
|
{
|
|
"id": 1,
|
|
"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._BeamlineConfig__client.delete.assert_called_with(ref_key)
|
|
|
|
|
|
def test_on_message_invalid_json(mock_config):
|
|
on_message(None, "invalid json")
|
|
mock_config._BeamlineConfig__client.set.assert_not_called() |