97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from control_backend.api.v1.endpoints import user_interact
|
|
|
|
|
|
@pytest.fixture
|
|
def app():
|
|
app = FastAPI()
|
|
app.include_router(user_interact.router)
|
|
return app
|
|
|
|
|
|
@pytest.fixture
|
|
def client(app):
|
|
return TestClient(app)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_receive_button_event(client):
|
|
mock_pub_socket = AsyncMock()
|
|
client.app.state.endpoints_pub_socket = mock_pub_socket
|
|
|
|
payload = {"type": "speech", "context": "hello"}
|
|
response = client.post("/button_pressed", json=payload)
|
|
|
|
assert response.status_code == 202
|
|
assert response.json() == {"status": "Event received"}
|
|
|
|
mock_pub_socket.send_multipart.assert_awaited_once()
|
|
args = mock_pub_socket.send_multipart.call_args[0][0]
|
|
assert args[0] == b"button_pressed"
|
|
assert "speech" in args[1].decode()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_receive_button_event_invalid_payload(client):
|
|
mock_pub_socket = AsyncMock()
|
|
client.app.state.endpoints_pub_socket = mock_pub_socket
|
|
|
|
# Missing context
|
|
payload = {"type": "speech"}
|
|
response = client.post("/button_pressed", json=payload)
|
|
|
|
assert response.status_code == 422
|
|
mock_pub_socket.send_multipart.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_experiment_stream_direct_call():
|
|
"""
|
|
Directly calling the endpoint function to test the streaming logic
|
|
without dealing with TestClient streaming limitations.
|
|
"""
|
|
mock_socket = AsyncMock()
|
|
# 1. recv data
|
|
# 2. recv timeout
|
|
# 3. disconnect (request.is_disconnected returns True)
|
|
mock_socket.recv_multipart.side_effect = [
|
|
(b"topic", b"message1"),
|
|
TimeoutError(),
|
|
(b"topic", b"message2"), # Should not be reached if disconnect checks work
|
|
]
|
|
mock_socket.close = MagicMock()
|
|
mock_socket.connect = MagicMock()
|
|
mock_socket.subscribe = MagicMock()
|
|
|
|
mock_context = MagicMock()
|
|
mock_context.socket.return_value = mock_socket
|
|
|
|
with patch(
|
|
"control_backend.api.v1.endpoints.user_interact.Context.instance", return_value=mock_context
|
|
):
|
|
mock_request = AsyncMock()
|
|
# is_disconnected sequence:
|
|
# 1. False (before first recv) -> reads message1
|
|
# 2. False (before second recv) -> triggers TimeoutError, continues
|
|
# 3. True (before third recv) -> break loop
|
|
mock_request.is_disconnected.side_effect = [False, False, True]
|
|
|
|
response = await user_interact.experiment_stream(mock_request)
|
|
|
|
lines = []
|
|
# Consume the generator
|
|
async for line in response.body_iterator:
|
|
lines.append(line)
|
|
|
|
assert "data: message1\n\n" in lines
|
|
assert len(lines) == 1
|
|
|
|
mock_socket.connect.assert_called()
|
|
mock_socket.subscribe.assert_called_with(b"experiment")
|
|
mock_socket.close.assert_called()
|