277 lines
12 KiB
Python
277 lines
12 KiB
Python
"""Module to test the Timepix Fly client functionality."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from ophyd import StatusBase
|
|
from websockets import State
|
|
|
|
from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client import (
|
|
TimepixFlyClient,
|
|
TimePixFlyStatus,
|
|
TimePixStatusError,
|
|
)
|
|
from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface import ProgramState
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def timepix_fly_client():
|
|
"""Fixture for creating a Timepix Fly client instance."""
|
|
client = TimepixFlyClient(rest_url="http://localhost:8000", ws_url="ws://localhost:8000/ws")
|
|
try:
|
|
yield client
|
|
finally:
|
|
client.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"return_state",
|
|
[
|
|
ProgramState(state=TimePixFlyStatus.CONFIG),
|
|
ProgramState(state=TimePixFlyStatus.COLLECT),
|
|
ProgramState(state=TimePixFlyStatus.SETUP),
|
|
],
|
|
)
|
|
def test_timepix_fly_client_stop_running_collection(timepix_fly_client, return_state):
|
|
"""Test the on_connected method of the Timepix Fly client."""
|
|
with (
|
|
mock.patch.object(timepix_fly_client, "stop_collect") as mock_stop_collect,
|
|
mock.patch.object(timepix_fly_client, "state", return_value=return_state),
|
|
):
|
|
timepix_fly_client.stop_running_collection()
|
|
|
|
if return_state.state == TimePixFlyStatus.CONFIG:
|
|
assert mock_stop_collect.call_count == 0, "Stop collect should be called once."
|
|
timepix_fly_client._started = True
|
|
timepix_fly_client.stop_running_collection()
|
|
assert mock_stop_collect.call_count == 1, "Stop collect should not be called again."
|
|
else:
|
|
assert mock_stop_collect.call_count == 1, "Stop collect should be called once."
|
|
timepix_fly_client._started = True
|
|
timepix_fly_client.stop_running_collection()
|
|
assert mock_stop_collect.call_count == 2, "Stop collect should be called in CONFIG."
|
|
|
|
|
|
def test_timepix_fly_client_on_connected(timepix_fly_client):
|
|
"""
|
|
Test timepix fly client connect method.
|
|
|
|
This simply ensures that all methods are called. They are tested separately.
|
|
"""
|
|
with (
|
|
mock.patch.object(timepix_fly_client, "stop_running_collection") as mock_stop_collection,
|
|
mock.patch.object(timepix_fly_client, "connect") as mock_connect,
|
|
mock.patch.object(timepix_fly_client, "wait_for_connection") as mock_wait_for_connection,
|
|
):
|
|
timepix_fly_client.on_connected()
|
|
mock_stop_collection.assert_called_once()
|
|
mock_connect.assert_called_once()
|
|
mock_wait_for_connection.assert_called_once()
|
|
|
|
|
|
def test_timepix_fly_client_connect(timepix_fly_client):
|
|
"""This tests the connect method of timepix fly client."""
|
|
module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
|
|
# Patch Thread so we don't actually start the background loop
|
|
with mock.patch(f"{module_path}.threading.Thread") as mock_thread_cls:
|
|
mock_thread = mock.Mock()
|
|
mock_thread_cls.return_value = mock_thread
|
|
|
|
timepix_fly_client.connect()
|
|
|
|
# Thread should be created with the update loop as target and daemon True
|
|
mock_thread_cls.assert_called_once()
|
|
# start() must be called on the created thread
|
|
mock_thread.start.assert_called_once()
|
|
|
|
|
|
def test_timepix_fly_client_wait_for_connection(timepix_fly_client):
|
|
"""This tests the wait_for_connection method of timepix fly client."""
|
|
module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
|
|
# Case 1: ws_client already present and OPEN
|
|
mock_client = mock.Mock()
|
|
type(mock_client).state = mock.PropertyMock(return_value=State.OPEN)
|
|
timepix_fly_client.ws_client = mock_client
|
|
|
|
# Should return immediately
|
|
timepix_fly_client.wait_for_connection(timeout=0.1)
|
|
|
|
# Case 2: connect() establishes the connection
|
|
timepix_fly_client.ws_client = None
|
|
mock_client2 = mock.Mock()
|
|
type(mock_client2).state = mock.PropertyMock(return_value=State.OPEN)
|
|
with mock.patch(f"{module_path}.connect", return_value=mock_client2) as mock_connect:
|
|
timepix_fly_client.wait_for_connection(timeout=0.1)
|
|
mock_connect.assert_called_once()
|
|
|
|
|
|
def test_timepix_fly_client_ws_send_and_received(timepix_fly_client):
|
|
"""This tests the _ws_send_and_receive method of timepix fly client."""
|
|
# Prepare a mock ws_client where recv returns a message
|
|
mock_ws = mock.Mock()
|
|
mock_ws.recv.return_value = "init"
|
|
timepix_fly_client.ws_client = mock_ws
|
|
|
|
# Patch _on_received_ws_message to ensure it's invoked
|
|
with mock.patch.object(timepix_fly_client, "_on_received_ws_message") as mock_on_msg:
|
|
timepix_fly_client._ws_send_and_receive()
|
|
mock_on_msg.assert_called_once_with("init")
|
|
|
|
# Now simulate a TimeoutError from recv: should simply return and not call _on_received_ws_message
|
|
mock_ws2 = mock.Mock()
|
|
mock_ws2.recv.side_effect = TimeoutError
|
|
timepix_fly_client.ws_client = mock_ws2
|
|
with mock.patch.object(timepix_fly_client, "_on_received_ws_message") as mock_on_msg2:
|
|
timepix_fly_client._ws_send_and_receive()
|
|
mock_on_msg2.assert_not_called()
|
|
|
|
|
|
def test_timepix_fly_client_on_message_received(timepix_fly_client):
|
|
"""This tests the _on_received_ws_message method of timepix fly client."""
|
|
with mock.patch.object(timepix_fly_client, "_run_status_callbacks") as mock_run_callbacks:
|
|
timepix_fly_client._on_received_ws_message("init")
|
|
assert timepix_fly_client._status == TimePixFlyStatus.INIT
|
|
mock_run_callbacks.assert_called_once()
|
|
|
|
# invalid message should not change status or call callbacks
|
|
prev_status = timepix_fly_client._status
|
|
with mock.patch.object(timepix_fly_client, "_run_status_callbacks") as mock_run_callbacks2:
|
|
timepix_fly_client._on_received_ws_message("invalid_status_string")
|
|
# status stays as previous value and callbacks not called
|
|
assert timepix_fly_client._status == prev_status
|
|
mock_run_callbacks2.assert_not_called()
|
|
|
|
|
|
def test_timepix_fly_client_on_status_callbacks(timepix_fly_client):
|
|
"""This tests the _run_status_callbacks method of timepix fly client."""
|
|
# Immediate run when current status already in success
|
|
timepix_fly_client._status = TimePixFlyStatus.INIT
|
|
status = StatusBase()
|
|
timepix_fly_client.add_status_callback(
|
|
status=status, success=[TimePixFlyStatus.INIT], error=[TimePixFlyStatus.EXCEPT], run=True
|
|
)
|
|
assert status.done is True and status.success is True
|
|
|
|
# Add callback (do not run immediately) and then trigger via _run_status_callbacks
|
|
status2 = StatusBase()
|
|
timepix_fly_client.add_status_callback(
|
|
status=status2,
|
|
success=[TimePixFlyStatus.CONFIG],
|
|
error=[TimePixFlyStatus.EXCEPT],
|
|
run=False,
|
|
)
|
|
# Set status to CONFIG and mark started True to check reset
|
|
timepix_fly_client._status = TimePixFlyStatus.CONFIG
|
|
timepix_fly_client._started = True
|
|
timepix_fly_client._run_status_callbacks()
|
|
assert status2.done is True and status2.success is True
|
|
assert timepix_fly_client._started is False
|
|
|
|
# Error path: add with run True when status is EXCEPT
|
|
timepix_fly_client._status = TimePixFlyStatus.EXCEPT
|
|
status3 = StatusBase()
|
|
with mock.patch.object(timepix_fly_client, "last_error") as mock_last_error:
|
|
mock_err = mock.Mock()
|
|
mock_err.message = "boom"
|
|
mock_last_error.return_value = mock_err
|
|
|
|
timepix_fly_client.add_status_callback(
|
|
status=status3,
|
|
success=[TimePixFlyStatus.INIT],
|
|
error=[TimePixFlyStatus.EXCEPT],
|
|
run=True,
|
|
)
|
|
assert status3.done is True and status3.success is False
|
|
|
|
|
|
def test_timepix_fly_client_shutdown(timepix_fly_client):
|
|
"""This tests the shutdown method of timepix fly client."""
|
|
mock_ws = mock.Mock()
|
|
timepix_fly_client.ws_client = mock_ws
|
|
|
|
timepix_fly_client.shutdown()
|
|
|
|
mock_ws.close.assert_called_once()
|
|
assert timepix_fly_client.ws_client is None
|
|
assert timepix_fly_client._shutdown_event.is_set()
|
|
|
|
|
|
def test_timepix_fly_client_start(timepix_fly_client):
|
|
"""This tests the start method of timepix fly client."""
|
|
# The client._get uses requests.get with f"http://{self.rest_url}/{get_cmd}".
|
|
# We mock requests.get and verify it is called with the expected URL and timeout,
|
|
# and that the _started flag is set.
|
|
module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
|
|
with mock.patch(f"{module_path}.requests.get") as mock_get:
|
|
mock_resp = mock.Mock()
|
|
mock_resp.raise_for_status = mock.Mock()
|
|
mock_resp.text = ""
|
|
mock_get.return_value = mock_resp
|
|
|
|
timepix_fly_client.start()
|
|
|
|
expected_url = f"http://{timepix_fly_client.rest_url}/?start=true"
|
|
mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout)
|
|
assert timepix_fly_client._started is True
|
|
|
|
|
|
def test_timepix_fly_client_stop_collect(timepix_fly_client):
|
|
"""This tests the stop_collect method of timepix fly client."""
|
|
module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
|
|
with mock.patch(f"{module_path}.requests.get") as mock_get:
|
|
mock_resp = mock.Mock()
|
|
mock_resp.raise_for_status = mock.Mock()
|
|
mock_resp.text = ""
|
|
mock_get.return_value = mock_resp
|
|
|
|
# ensure started is True and then call stop_collect
|
|
timepix_fly_client._started = True
|
|
timepix_fly_client.stop_collect()
|
|
|
|
expected_url = f"http://{timepix_fly_client.rest_url}/?stop_collect=true"
|
|
mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout)
|
|
assert timepix_fly_client._started is False
|
|
|
|
|
|
def test_timepix_fly_client_state(timepix_fly_client):
|
|
"""This tests the state method of timepix fly client."""
|
|
module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
|
|
with mock.patch(f"{module_path}.requests.get") as mock_get:
|
|
mock_resp = mock.Mock()
|
|
mock_resp.raise_for_status = mock.Mock()
|
|
# Return a JSON payload compatible with ProgramState
|
|
mock_resp.json.return_value = {"type": "ProgramState", "state": "init"}
|
|
mock_get.return_value = mock_resp
|
|
|
|
program_state = timepix_fly_client.state()
|
|
|
|
expected_url = f"http://{timepix_fly_client.rest_url}/state"
|
|
mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout)
|
|
# ProgramState.model defines 'state' as a string literal; ensure we parsed it
|
|
assert hasattr(program_state, "state")
|
|
assert program_state.state == "init"
|
|
|
|
|
|
def test_timepix_fly_client_set_pixel_map(timepix_fly_client):
|
|
"""This tests the set_pixel_map/_put path by mocking requests.put and checking payload."""
|
|
module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client"
|
|
with mock.patch(f"{module_path}.requests.put") as mock_put:
|
|
mock_resp = mock.Mock()
|
|
mock_resp.raise_for_status = mock.Mock()
|
|
mock_put.return_value = mock_resp
|
|
|
|
# Minimal valid PixelMap dict (type is optional; model supplies default)
|
|
pixel_map = {"chips": [[{"i": 0, "p": 1, "f": 2}]]}
|
|
|
|
timepix_fly_client.set_pixel_map(pixel_map)
|
|
|
|
expected_url = f"http://{timepix_fly_client.rest_url}/pixel-map"
|
|
# Verify requests.put called with expected url, json and timeout
|
|
mock_put.assert_called_once()
|
|
_, kwargs = mock_put.call_args
|
|
assert kwargs.get("timeout") == timepix_fly_client._timeout
|
|
assert kwargs.get("json") is not None
|
|
assert kwargs.get("json").get("chips") == pixel_map["chips"]
|