521 lines
16 KiB
Python
521 lines
16 KiB
Python
import json
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
import zmq
|
|
|
|
from control_backend.agents.actuation.robot_gesture_agent import RobotGestureAgent
|
|
from control_backend.core.agent_system import InternalMessage
|
|
from control_backend.schemas.ri_message import RIEndpoint
|
|
|
|
|
|
@pytest.fixture
|
|
def zmq_context(mocker):
|
|
"""Mock the ZMQ context."""
|
|
mock_context = mocker.patch(
|
|
"control_backend.agents.actuation.robot_gesture_agent.azmq.Context.instance"
|
|
)
|
|
mock_context.return_value = MagicMock()
|
|
return mock_context
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_experiment_logger():
|
|
with patch("control_backend.agents.actuation.robot_gesture_agent.experiment_logger") as logger:
|
|
yield logger
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_setup_bind(zmq_context, mocker):
|
|
"""Setup binds and subscribes to internal commands."""
|
|
fake_socket = zmq_context.return_value.socket.return_value
|
|
agent = RobotGestureAgent("robot_gesture", address="tcp://localhost:5556", bind=True)
|
|
|
|
settings = mocker.patch("control_backend.agents.actuation.robot_gesture_agent.settings")
|
|
settings.zmq_settings.internal_sub_address = "tcp://internal:1234"
|
|
|
|
def close_coro(coro):
|
|
coro.close()
|
|
return MagicMock()
|
|
|
|
agent.add_behavior = MagicMock(side_effect=close_coro)
|
|
|
|
await agent.setup()
|
|
|
|
# Check PUB socket binding
|
|
fake_socket.bind.assert_any_call("tcp://localhost:5556")
|
|
# Check REP socket binding
|
|
fake_socket.bind.assert_called()
|
|
|
|
# Check SUB socket connection and subscriptions
|
|
fake_socket.connect.assert_any_call("tcp://internal:1234")
|
|
fake_socket.setsockopt.assert_any_call(zmq.SUBSCRIBE, b"command")
|
|
fake_socket.setsockopt.assert_any_call(zmq.SUBSCRIBE, b"send_gestures")
|
|
|
|
# Check behavior was added (twice: once for command loop, once for fetch gestures loop)
|
|
assert agent.add_behavior.call_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_setup_connect(zmq_context, mocker):
|
|
"""Setup connects when bind=False."""
|
|
fake_socket = zmq_context.return_value.socket.return_value
|
|
agent = RobotGestureAgent("robot_gesture", address="tcp://localhost:5556", bind=False)
|
|
|
|
settings = mocker.patch("control_backend.agents.actuation.robot_gesture_agent.settings")
|
|
settings.zmq_settings.internal_sub_address = "tcp://internal:1234"
|
|
|
|
def close_coro(coro):
|
|
coro.close()
|
|
return MagicMock()
|
|
|
|
agent.add_behavior = MagicMock(side_effect=close_coro)
|
|
|
|
await agent.setup()
|
|
|
|
# Check PUB socket connection (not binding)
|
|
fake_socket.connect.assert_any_call("tcp://localhost:5556")
|
|
fake_socket.connect.assert_any_call("tcp://internal:1234")
|
|
# Check REP socket binding (always binds)
|
|
fake_socket.bind.assert_called()
|
|
|
|
# Check behavior was added (twice)
|
|
assert agent.add_behavior.call_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_sends_valid_gesture_command():
|
|
"""Internal message with valid gesture tag is forwarded to robot pub socket."""
|
|
pubsocket = AsyncMock()
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.pubsocket = pubsocket
|
|
|
|
payload = {
|
|
"endpoint": RIEndpoint.GESTURE_TAG,
|
|
"data": "hello", # "hello" is in gesture_data
|
|
}
|
|
msg = InternalMessage(to="robot", sender="tester", body=json.dumps(payload))
|
|
|
|
await agent.handle_message(msg)
|
|
|
|
pubsocket.send_json.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_sends_non_gesture_command():
|
|
"""Internal message with non-gesture endpoint is not forwarded by this agent."""
|
|
pubsocket = AsyncMock()
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.pubsocket = pubsocket
|
|
|
|
payload = {"endpoint": "some_other_endpoint", "data": "invalid_tag_not_in_list"}
|
|
msg = InternalMessage(to="robot", sender="tester", body=json.dumps(payload))
|
|
|
|
await agent.handle_message(msg)
|
|
|
|
# Non-gesture endpoints should not be forwarded by this agent
|
|
pubsocket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_rejects_invalid_gesture_tag():
|
|
"""Internal message with invalid gesture tag is not forwarded."""
|
|
pubsocket = AsyncMock()
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.pubsocket = pubsocket
|
|
|
|
# Use a tag that's not in gesture_data
|
|
payload = {"endpoint": RIEndpoint.GESTURE_TAG, "data": "invalid_tag_not_in_list"}
|
|
msg = InternalMessage(to="robot", sender="tester", body=json.dumps(payload))
|
|
|
|
await agent.handle_message(msg)
|
|
|
|
pubsocket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_sends_valid_single_gesture_command():
|
|
"""Internal message with valid single gesture is forwarded."""
|
|
pubsocket = AsyncMock()
|
|
agent = RobotGestureAgent("robot_gesture", single_gesture_data=["wave", "point"], address="")
|
|
agent.pubsocket = pubsocket
|
|
|
|
payload = {
|
|
"endpoint": RIEndpoint.GESTURE_SINGLE,
|
|
"data": "wave",
|
|
}
|
|
msg = InternalMessage(to="robot", sender="tester", body=json.dumps(payload))
|
|
|
|
await agent.handle_message(msg)
|
|
|
|
pubsocket.send_json.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_rejects_invalid_single_gesture():
|
|
"""Internal message with invalid single gesture is not forwarded."""
|
|
pubsocket = AsyncMock()
|
|
agent = RobotGestureAgent("robot_gesture", single_gesture_data=["wave", "point"], address="")
|
|
agent.pubsocket = pubsocket
|
|
|
|
payload = {
|
|
"endpoint": RIEndpoint.GESTURE_SINGLE,
|
|
"data": "dance",
|
|
}
|
|
msg = InternalMessage(to="robot", sender="tester", body=json.dumps(payload))
|
|
|
|
await agent.handle_message(msg)
|
|
|
|
pubsocket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zmq_command_loop_valid_single_gesture_payload():
|
|
"""UI command with valid single gesture is read from SUB and published."""
|
|
command = {"endpoint": RIEndpoint.GESTURE_SINGLE, "data": "wave"}
|
|
fake_socket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return b"command", json.dumps(command).encode("utf-8")
|
|
|
|
fake_socket.recv_multipart = recv_once
|
|
fake_socket.send_json = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", single_gesture_data=["wave", "point"], address="")
|
|
agent.subsocket = fake_socket
|
|
agent.pubsocket = fake_socket
|
|
agent._running = True
|
|
|
|
await agent._zmq_command_loop()
|
|
|
|
fake_socket.send_json.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_message_invalid_payload():
|
|
"""Invalid payload is caught and does not send."""
|
|
pubsocket = AsyncMock()
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.pubsocket = pubsocket
|
|
|
|
msg = InternalMessage(to="robot", sender="tester", body=json.dumps({"bad": "data"}))
|
|
|
|
await agent.handle_message(msg)
|
|
|
|
pubsocket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zmq_command_loop_valid_gesture_payload():
|
|
"""UI command with valid gesture tag is read from SUB and published."""
|
|
command = {"endpoint": RIEndpoint.GESTURE_TAG, "data": "hello"}
|
|
fake_socket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
# stop after first iteration
|
|
agent._running = False
|
|
return b"command", json.dumps(command).encode("utf-8")
|
|
|
|
fake_socket.recv_multipart = recv_once
|
|
fake_socket.send_json = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.subsocket = fake_socket
|
|
agent.pubsocket = fake_socket
|
|
agent._running = True
|
|
|
|
await agent._zmq_command_loop()
|
|
|
|
fake_socket.send_json.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zmq_command_loop_valid_non_gesture_payload():
|
|
"""UI command with non-gesture endpoint is not forwarded by this agent."""
|
|
command = {"endpoint": "some_other_endpoint", "data": "anything"}
|
|
fake_socket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return b"command", json.dumps(command).encode("utf-8")
|
|
|
|
fake_socket.recv_multipart = recv_once
|
|
fake_socket.send_json = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.subsocket = fake_socket
|
|
agent.pubsocket = fake_socket
|
|
agent._running = True
|
|
|
|
await agent._zmq_command_loop()
|
|
|
|
fake_socket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zmq_command_loop_invalid_gesture_tag():
|
|
"""UI command with invalid gesture tag is not forwarded."""
|
|
command = {"endpoint": RIEndpoint.GESTURE_TAG, "data": "invalid_tag_not_in_list"}
|
|
fake_socket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return b"command", json.dumps(command).encode("utf-8")
|
|
|
|
fake_socket.recv_multipart = recv_once
|
|
fake_socket.send_json = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.subsocket = fake_socket
|
|
agent.pubsocket = fake_socket
|
|
agent._running = True
|
|
|
|
await agent._zmq_command_loop()
|
|
|
|
fake_socket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zmq_command_loop_invalid_json():
|
|
"""Invalid JSON is ignored without sending."""
|
|
fake_socket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return b"command", b"{not_json}"
|
|
|
|
fake_socket.recv_multipart = recv_once
|
|
fake_socket.send_json = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.subsocket = fake_socket
|
|
agent.pubsocket = fake_socket
|
|
agent._running = True
|
|
|
|
await agent._zmq_command_loop()
|
|
|
|
fake_socket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zmq_command_loop_ignores_send_gestures_topic():
|
|
"""send_gestures topic is ignored in command loop."""
|
|
fake_socket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return b"send_gestures", b"{}"
|
|
|
|
fake_socket.recv_multipart = recv_once
|
|
fake_socket.send_json = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.subsocket = fake_socket
|
|
agent.pubsocket = fake_socket
|
|
agent._running = True
|
|
|
|
await agent._zmq_command_loop()
|
|
|
|
fake_socket.send_json.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_gestures_loop_without_amount():
|
|
"""Fetch gestures request without amount returns all tags."""
|
|
fake_repsocket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return b"{}" # Empty JSON request
|
|
|
|
fake_repsocket.recv = recv_once
|
|
fake_repsocket.send = AsyncMock()
|
|
|
|
agent = RobotGestureAgent(
|
|
"robot_gesture", gesture_data=["hello", "yes", "no", "wave", "point"], address=""
|
|
)
|
|
agent.repsocket = fake_repsocket
|
|
agent._running = True
|
|
|
|
await agent._fetch_gestures_loop()
|
|
|
|
fake_repsocket.send.assert_awaited_once()
|
|
|
|
# Check the response contains all tags
|
|
args, kwargs = fake_repsocket.send.call_args
|
|
response = json.loads(args[0])
|
|
assert "tags" in response
|
|
assert response["tags"] == ["hello", "yes", "no", "wave", "point"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_gestures_loop_with_amount():
|
|
"""Fetch gestures request with amount returns limited tags."""
|
|
fake_repsocket = AsyncMock()
|
|
amount = 3
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return json.dumps(amount).encode()
|
|
|
|
fake_repsocket.recv = recv_once
|
|
fake_repsocket.send = AsyncMock()
|
|
|
|
agent = RobotGestureAgent(
|
|
"robot_gesture", gesture_data=["hello", "yes", "no", "wave", "point"], address=""
|
|
)
|
|
agent.repsocket = fake_repsocket
|
|
agent._running = True
|
|
|
|
await agent._fetch_gestures_loop()
|
|
|
|
fake_repsocket.send.assert_awaited_once()
|
|
|
|
args, kwargs = fake_repsocket.send.call_args
|
|
response = json.loads(args[0])
|
|
assert "tags" in response
|
|
assert len(response["tags"]) == amount
|
|
assert response["tags"] == ["hello", "yes", "no"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_gestures_loop_with_integer_request():
|
|
"""Fetch gestures request with integer amount."""
|
|
fake_repsocket = AsyncMock()
|
|
amount = 2
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return json.dumps(amount).encode()
|
|
|
|
fake_repsocket.recv = recv_once
|
|
fake_repsocket.send = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.repsocket = fake_repsocket
|
|
agent._running = True
|
|
|
|
await agent._fetch_gestures_loop()
|
|
|
|
fake_repsocket.send.assert_awaited_once()
|
|
|
|
args, kwargs = fake_repsocket.send.call_args
|
|
response = json.loads(args[0])
|
|
assert response["tags"] == ["hello", "yes"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_gestures_loop_with_invalid_json():
|
|
"""Invalid JSON request returns all tags."""
|
|
fake_repsocket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return b"not_json"
|
|
|
|
fake_repsocket.recv = recv_once
|
|
fake_repsocket.send = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.repsocket = fake_repsocket
|
|
agent._running = True
|
|
|
|
await agent._fetch_gestures_loop()
|
|
|
|
fake_repsocket.send.assert_awaited_once()
|
|
|
|
args, kwargs = fake_repsocket.send.call_args
|
|
response = json.loads(args[0])
|
|
assert response["tags"] == ["hello", "yes", "no"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_gestures_loop_with_non_integer_json():
|
|
"""Non-integer JSON request returns all tags."""
|
|
fake_repsocket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
return json.dumps({"not": "an_integer"}).encode()
|
|
|
|
fake_repsocket.recv = recv_once
|
|
fake_repsocket.send = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.repsocket = fake_repsocket
|
|
agent._running = True
|
|
|
|
await agent._fetch_gestures_loop()
|
|
|
|
fake_repsocket.send.assert_awaited_once()
|
|
|
|
args, kwargs = fake_repsocket.send.call_args
|
|
response = json.loads(args[0])
|
|
assert response["tags"] == ["hello", "yes", "no"]
|
|
|
|
|
|
def test_gesture_data_attribute():
|
|
"""Test that gesture_data returns the expected list."""
|
|
gesture_data = ["hello", "yes", "no", "wave"]
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=gesture_data, address="")
|
|
|
|
assert agent.gesture_data == gesture_data
|
|
assert isinstance(agent.gesture_data, list)
|
|
assert len(agent.gesture_data) == 4
|
|
assert "hello" in agent.gesture_data
|
|
assert "yes" in agent.gesture_data
|
|
assert "no" in agent.gesture_data
|
|
assert "invalid_tag_not_in_list" not in agent.gesture_data
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_closes_sockets():
|
|
"""Stop method closes all sockets."""
|
|
pubsocket = MagicMock()
|
|
subsocket = MagicMock()
|
|
repsocket = MagicMock()
|
|
agent = RobotGestureAgent("robot_gesture", address="")
|
|
agent.pubsocket = pubsocket
|
|
agent.subsocket = subsocket
|
|
agent.repsocket = repsocket
|
|
|
|
await agent.stop()
|
|
|
|
pubsocket.close.assert_called_once()
|
|
subsocket.close.assert_called_once()
|
|
repsocket.close.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialization_with_custom_gesture_data():
|
|
"""Agent can be initialized with custom gesture data."""
|
|
custom_gestures = ["custom1", "custom2", "custom3"]
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=custom_gestures, address="")
|
|
|
|
assert agent.gesture_data == custom_gestures
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_gestures_loop_handles_exception():
|
|
"""Exception in fetch gestures loop is caught and logged."""
|
|
fake_repsocket = AsyncMock()
|
|
|
|
async def recv_once():
|
|
agent._running = False
|
|
raise Exception("Test exception")
|
|
|
|
fake_repsocket.recv = recv_once
|
|
fake_repsocket.send = AsyncMock()
|
|
|
|
agent = RobotGestureAgent("robot_gesture", gesture_data=["hello", "yes", "no"], address="")
|
|
agent.repsocket = fake_repsocket
|
|
agent.logger = MagicMock()
|
|
agent._running = True
|
|
|
|
# Should not raise exception
|
|
await agent._fetch_gestures_loop()
|
|
|
|
# Exception should be logged
|
|
agent.logger.exception.assert_called_once()
|