chore: added missing tests
This commit is contained in:
@@ -94,3 +94,55 @@ async def test_experiment_stream_direct_call():
|
||||
mock_socket.connect.assert_called()
|
||||
mock_socket.subscribe.assert_called_with(b"experiment")
|
||||
mock_socket.close.assert_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_status_stream_direct_call():
|
||||
"""
|
||||
Test the status stream, ensuring it handles messages and sends pings on timeout.
|
||||
"""
|
||||
mock_socket = AsyncMock()
|
||||
|
||||
# Define the sequence of events for the socket:
|
||||
# 1. Successfully receive a message
|
||||
# 2. Timeout (which should trigger the ': ping' yield)
|
||||
# 3. Another message (which won't be reached because we'll simulate disconnect)
|
||||
mock_socket.recv_multipart.side_effect = [
|
||||
(b"topic", b"status_update"),
|
||||
TimeoutError(),
|
||||
(b"topic", b"ignored_msg"),
|
||||
]
|
||||
|
||||
mock_socket.close = MagicMock()
|
||||
mock_socket.connect = MagicMock()
|
||||
mock_socket.subscribe = MagicMock()
|
||||
|
||||
mock_context = MagicMock()
|
||||
mock_context.socket.return_value = mock_socket
|
||||
|
||||
# Mock the ZMQ Context to return our 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 -> Process "status_update"
|
||||
# 2. False -> Process TimeoutError (yield ping)
|
||||
# 3. True -> Break loop (client disconnected)
|
||||
mock_request.is_disconnected.side_effect = [False, False, True]
|
||||
|
||||
# Call the status_stream function explicitly
|
||||
response = await user_interact.status_stream(mock_request)
|
||||
|
||||
lines = []
|
||||
async for line in response.body_iterator:
|
||||
lines.append(line)
|
||||
|
||||
# Assertions
|
||||
assert "data: status_update\n\n" in lines
|
||||
assert ": ping\n\n" in lines # Verify lines 91-92 (ping logic)
|
||||
|
||||
mock_socket.connect.assert_called()
|
||||
mock_socket.subscribe.assert_called_with(b"status")
|
||||
mock_socket.close.assert_called()
|
||||
|
||||
Reference in New Issue
Block a user