Separate and tested web_socket_manager
This commit is contained in:
@@ -1,53 +0,0 @@
|
||||
from typing import Optional, List, Callable
|
||||
from sia.web_agent import WebAgentState
|
||||
|
||||
class MockWebAgent:
|
||||
"""Mock WebAgent for testing."""
|
||||
def __init__(self):
|
||||
self._current_state: Optional[WebAgentState] = None
|
||||
self.context: Optional[str] = None
|
||||
self.response: Optional[str] = None
|
||||
self._validation_error: Optional[str] = None
|
||||
self._state_handlers: List[Callable] = []
|
||||
self._response_handlers: List[Callable] = []
|
||||
|
||||
@property
|
||||
def current_state(self):
|
||||
return self._current_state
|
||||
|
||||
def add_state_change_handler(self, handler):
|
||||
self._state_handlers.append(handler)
|
||||
|
||||
def add_response_change_handler(self, handler):
|
||||
self._response_handlers.append(handler)
|
||||
|
||||
def set_input(self, input_text):
|
||||
pass
|
||||
|
||||
def _notify_state_handlers(self, new_state: WebAgentState):
|
||||
"""Notify all state handlers of a state change."""
|
||||
for handler in self._state_handlers:
|
||||
handler(new_state)
|
||||
|
||||
def _notify_response_handlers(self, response: str):
|
||||
"""Notify all response handlers of a response change."""
|
||||
for handler in self._response_handlers:
|
||||
handler(response)
|
||||
|
||||
def _set_response(self, response: str):
|
||||
"""Set response and notify handlers."""
|
||||
self.response = response
|
||||
self._notify_response_handlers(response)
|
||||
|
||||
def approve_context(self):
|
||||
"""Handle context approval with immediate state updates."""
|
||||
# Change to INFERENCE state and notify
|
||||
self._current_state = WebAgentState.INFERENCE
|
||||
self._notify_state_handlers(WebAgentState.INFERENCE)
|
||||
|
||||
# Set response and notify
|
||||
self._set_response("<reasoning>test</reasoning>")
|
||||
|
||||
# Change to RESPONSE_APPROVAL state and notify
|
||||
self._current_state = WebAgentState.RESPONSE_APPROVAL
|
||||
self._notify_state_handlers(WebAgentState.RESPONSE_APPROVAL)
|
||||
@@ -88,7 +88,7 @@ class WebAgentTest(unittest.TestCase):
|
||||
"""Track state changes for verification."""
|
||||
self.state_changes.append(state)
|
||||
|
||||
def response_change_handler(self, response: str):
|
||||
def response_change_handler(self, response: str, validation_error: str):
|
||||
"""Track response changes for verification."""
|
||||
self.response_changes.append(response)
|
||||
|
||||
@@ -97,7 +97,7 @@ class WebAgentTest(unittest.TestCase):
|
||||
self.assertEqual(self.agent.current_state, WebAgentState.CONTEXT_APPROVAL)
|
||||
self.assertEqual(self.agent.response, "")
|
||||
self.assertIsNotNone(self.agent.context)
|
||||
self.assertIsNone(self.agent._validation_error)
|
||||
self.assertIsNone(self.agent.validation_error)
|
||||
self.assertEqual(len(self.agent._state_change_handlers), 0)
|
||||
self.assertEqual(len(self.agent._response_change_handlers), 0)
|
||||
|
||||
@@ -133,7 +133,7 @@ class WebAgentTest(unittest.TestCase):
|
||||
self.agent.approve_response()
|
||||
self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
|
||||
|
||||
async def test_context_approval_flow(self):
|
||||
def test_context_approval_flow(self):
|
||||
"""Test complete context approval flow with state transitions."""
|
||||
# Register handlers
|
||||
self.agent.add_state_change_handler(self.state_change_handler)
|
||||
@@ -145,27 +145,16 @@ class WebAgentTest(unittest.TestCase):
|
||||
yield "<reasoning>test reasoning</reasoning>"
|
||||
self.mock_llm.infer.side_effect = mock_infer
|
||||
|
||||
# Set initial state
|
||||
self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
|
||||
|
||||
# Approve context
|
||||
self.agent.approve_context()
|
||||
|
||||
# Wait for response tasks to complete
|
||||
await asyncio.gather(*self.web_server._response_tasks)
|
||||
|
||||
# Verify state transitions
|
||||
expected_states = [
|
||||
WebAgentState.INFERENCE,
|
||||
WebAgentState.RESPONSE_APPROVAL
|
||||
]
|
||||
self.assertEqual(self.state_changes, expected_states)
|
||||
|
||||
# Updated test: Should now expect streaming updates
|
||||
self.assertEqual(len(self.response_changes), len(self.received_messages))
|
||||
for i, msg in enumerate(self.received_messages):
|
||||
self.assertEqual(msg["type"], "RESPONSE_UPDATE")
|
||||
self.assertEqual(msg["data"], self.response_changes[i])
|
||||
self.assertEqual(self.response_changes[-1], "<reasoning>test reasoning</reasoning>")
|
||||
|
||||
def test_response_approval_flow_command(self):
|
||||
"""Test response approval flow with command execution."""
|
||||
@@ -206,7 +195,7 @@ class WebAgentTest(unittest.TestCase):
|
||||
|
||||
# Set initial state and response
|
||||
self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
|
||||
self.agent.response = "<reasoning>test reasoning</reasoning>"
|
||||
self.agent.set_response("<reasoning>test reasoning</reasoning>")
|
||||
|
||||
# Approve response
|
||||
self.agent.approve_response()
|
||||
@@ -224,22 +213,13 @@ class WebAgentTest(unittest.TestCase):
|
||||
]
|
||||
self.assertEqual(self.state_changes, expected_states)
|
||||
|
||||
async def test_response_validation(self):
|
||||
def test_response_validation(self):
|
||||
"""Test response validation during response setting."""
|
||||
# Register handlers
|
||||
self.agent.add_response_change_handler(self.response_change_handler)
|
||||
|
||||
# Mock validator to fail
|
||||
error_message = "Invalid XML"
|
||||
self.mock_validator.validate.return_value = error_message
|
||||
|
||||
# Set response using async method
|
||||
await self.agent._set_response("<invalid>")
|
||||
|
||||
# Verify validation error stored
|
||||
self.assertEqual(self.agent._validation_error, error_message)
|
||||
|
||||
# Verify response updated
|
||||
self.agent.set_response("<invalid>")
|
||||
self.assertEqual(self.agent.validation_error, error_message)
|
||||
self.assertEqual(self.response_changes, ["<invalid>"])
|
||||
|
||||
def tearDown(self):
|
||||
|
||||
@@ -1,251 +0,0 @@
|
||||
from aiohttp import WSMsgType
|
||||
from aiohttp import WSMsgType
|
||||
from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
|
||||
from pathlib import Path
|
||||
from sia.web_agent import WebAgentState
|
||||
from unittest.mock import patch
|
||||
import asyncio
|
||||
import asyncio
|
||||
import time
|
||||
import json
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from sia.web_agent import WebAgentState
|
||||
from sia.web_io_buffer import WebIOBuffer
|
||||
from sia.web_server import WebServer, ServerMessage, ClientMessage
|
||||
from .mock_web_agent import MockWebAgent
|
||||
|
||||
class WebServerTest(AioHTTPTestCase):
|
||||
async def get_application(self):
|
||||
"""Create application for testing."""
|
||||
self.io_buffer = WebIOBuffer()
|
||||
|
||||
self.agent = MockWebAgent()
|
||||
|
||||
# Create temporary directory for static files
|
||||
self.temp_dir = Path(tempfile.mkdtemp())
|
||||
self.static_dir = self.temp_dir / "static"
|
||||
self.static_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Create web server with temp static directory
|
||||
self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
|
||||
self.web_server = WebServer(
|
||||
self.agent,
|
||||
self.io_buffer,
|
||||
"localhost",
|
||||
8080,
|
||||
static_dir=self.static_dir
|
||||
)
|
||||
return self.web_server.app
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up temporary directory."""
|
||||
if hasattr(self, 'temp_dir'):
|
||||
shutil.rmtree(self.temp_dir)
|
||||
|
||||
async def wait_for_messages(self, ws, count=1, timeout=5):
|
||||
"""Helper to wait for multiple WebSocket messages."""
|
||||
messages = []
|
||||
try:
|
||||
for _ in range(count):
|
||||
msg = await ws.receive(timeout=timeout)
|
||||
if msg.type == WSMsgType.TEXT:
|
||||
messages.append(json.loads(msg.data))
|
||||
elif msg.type == WSMsgType.BINARY:
|
||||
messages.append(json.loads(msg.data.decode()))
|
||||
elif msg.type == WSMsgType.CLOSED:
|
||||
break
|
||||
elif msg.type == WSMsgType.ERROR:
|
||||
self.fail(f"WebSocket error while waiting for messages: {msg.data}")
|
||||
return messages
|
||||
except asyncio.TimeoutError:
|
||||
self.fail(f"Timeout waiting for WebSocket message. Got {len(messages)} of {count}")
|
||||
except json.JSONDecodeError as e:
|
||||
self.fail(f"Invalid JSON in message: {e}")
|
||||
|
||||
async def drain_messages(self, ws, timeout=0.1):
|
||||
"""Helper to drain any pending messages."""
|
||||
try:
|
||||
while True:
|
||||
msg = await ws.receive(timeout=timeout)
|
||||
if msg.type == WSMsgType.CLOSED:
|
||||
break # Normal close, no need to fail
|
||||
elif msg.type == WSMsgType.ERROR:
|
||||
self.fail(f"WebSocket error while draining messages: {msg.data}")
|
||||
except asyncio.TimeoutError:
|
||||
pass # Expected when no more messages
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test case with initial state."""
|
||||
super().setUp()
|
||||
self.io_buffer = None
|
||||
self.agent = None
|
||||
self.web_server = None
|
||||
|
||||
async def get_application(self):
|
||||
"""Create application for testing."""
|
||||
self.io_buffer = WebIOBuffer()
|
||||
self.agent = MockWebAgent()
|
||||
self.agent._current_state = WebAgentState.CONTEXT_APPROVAL # Set initial state
|
||||
self.web_server = WebServer(self.agent, self.io_buffer, "localhost", 8080)
|
||||
return self.web_server.app
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_approve_context(self):
|
||||
"""Test context approval flow."""
|
||||
async with self.client.ws_connect("/ws") as ws:
|
||||
# Drain initial messages
|
||||
await self.drain_messages(ws)
|
||||
await asyncio.sleep(0.1) # Let connection settle
|
||||
|
||||
# Set agent to CONTEXT_APPROVAL state
|
||||
self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
|
||||
self.agent.context = "<test>context</test>"
|
||||
await asyncio.sleep(0.1) # Let state changes propagate
|
||||
|
||||
# Send context approval
|
||||
await ws.send_json({
|
||||
"type": ClientMessage.APPROVE_CONTEXT,
|
||||
})
|
||||
await asyncio.sleep(0.1) # Let state changes propagate
|
||||
|
||||
# Should receive 3 messages: INFERENCE state, response, RESPONSE_APPROVAL state
|
||||
messages = await self.wait_for_messages(ws, count=3)
|
||||
|
||||
# Verify state change to INFERENCE
|
||||
self.assertEqual(messages[0]["type"], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(messages[0]["data"], WebAgentState.INFERENCE.name)
|
||||
|
||||
# Verify response
|
||||
self.assertEqual(messages[1]["type"], ServerMessage.RESPONSE_UPDATE)
|
||||
self.assertEqual(messages[1]["data"], "<reasoning>test</reasoning>")
|
||||
|
||||
# Verify state change to RESPONSE_APPROVAL
|
||||
self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name)
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_validation_error(self):
|
||||
"""Test handling of validation errors."""
|
||||
async with self.client.ws_connect("/ws") as ws:
|
||||
# Drain initial messages
|
||||
await self.drain_messages(ws)
|
||||
await asyncio.sleep(0.1) # Let connection settle
|
||||
|
||||
# Set validation error and response
|
||||
self.agent._validation_error = "Invalid XML"
|
||||
self.agent._set_response("<invalid>") # Now synchronous
|
||||
await asyncio.sleep(0.1) # Let changes propagate
|
||||
|
||||
# Should receive response update with error
|
||||
messages = await self.wait_for_messages(ws, count=1)
|
||||
self.assertEqual(messages[0]["type"], ServerMessage.RESPONSE_UPDATE)
|
||||
self.assertEqual(messages[0]["validation_error"], "Invalid XML")
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_multiple_clients(self):
|
||||
"""Test handling multiple client connections."""
|
||||
ws1 = await self.client.ws_connect("/ws")
|
||||
ws2 = await self.client.ws_connect("/ws")
|
||||
|
||||
try:
|
||||
# Drain initial messages
|
||||
await self.drain_messages(ws1)
|
||||
await self.drain_messages(ws2)
|
||||
await asyncio.sleep(0.1) # Let connections settle
|
||||
|
||||
# Verify both clients are tracked
|
||||
self.assertEqual(len(self.web_server.clients), 2)
|
||||
|
||||
finally:
|
||||
await ws1.close()
|
||||
await ws2.close()
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_client_disconnect(self):
|
||||
"""Test proper cleanup on client disconnect."""
|
||||
ws = await self.client.ws_connect("/ws")
|
||||
await self.drain_messages(ws)
|
||||
await asyncio.sleep(0.1) # Let connection settle
|
||||
|
||||
self.assertEqual(len(self.web_server.clients), 1)
|
||||
|
||||
await ws.close()
|
||||
await asyncio.sleep(0.1) # Let cleanup complete
|
||||
|
||||
self.assertEqual(len(self.web_server.clients), 0)
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_broadcast(self):
|
||||
"""Test broadcasting to all clients."""
|
||||
ws1 = await self.client.ws_connect("/ws")
|
||||
ws2 = await self.client.ws_connect("/ws")
|
||||
|
||||
try:
|
||||
# Drain initial messages
|
||||
await self.drain_messages(ws1)
|
||||
await self.drain_messages(ws2)
|
||||
await asyncio.sleep(0.1) # Let connections settle
|
||||
|
||||
# Broadcast test message
|
||||
test_msg = "test broadcast"
|
||||
await self.web_server.broadcast(ServerMessage.OUTPUT_UPDATE, test_msg)
|
||||
|
||||
# Both clients should receive message
|
||||
msg1 = (await self.wait_for_messages(ws1, count=1))[0]
|
||||
msg2 = (await self.wait_for_messages(ws2, count=1))[0]
|
||||
|
||||
self.assertEqual(msg1["type"], ServerMessage.OUTPUT_UPDATE)
|
||||
self.assertEqual(msg1["data"], test_msg)
|
||||
self.assertEqual(msg2["type"], ServerMessage.OUTPUT_UPDATE)
|
||||
self.assertEqual(msg2["data"], test_msg)
|
||||
|
||||
finally:
|
||||
await ws1.close()
|
||||
await ws2.close()
|
||||
|
||||
async def wait_for_buffer_content(self, timeout=1.0):
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
if self.io_buffer.buffer_length() > 0:
|
||||
return True
|
||||
await asyncio.sleep(0.01)
|
||||
return False
|
||||
|
||||
async def test_send_input(self):
|
||||
async with self.client.ws_connect("/ws") as ws:
|
||||
await self.drain_messages(ws)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
test_input = "test input"
|
||||
await ws.send_json({
|
||||
"type": ClientMessage.SEND_INPUT,
|
||||
"data": test_input
|
||||
})
|
||||
|
||||
# Wait for buffer to have content
|
||||
self.assertTrue(await self.wait_for_buffer_content())
|
||||
|
||||
content = self.io_buffer.read()
|
||||
self.assertEqual(
|
||||
content,
|
||||
test_input + "\n",
|
||||
f"Expected '{test_input}\\n' but got '{content}'"
|
||||
)
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_invalid_json(self):
|
||||
"""Test handling of invalid JSON messages."""
|
||||
async with self.client.ws_connect("/ws") as ws:
|
||||
# Drain initial messages
|
||||
await self.drain_messages(ws)
|
||||
await asyncio.sleep(0.1) # Let connection settle
|
||||
|
||||
# Send invalid JSON
|
||||
await ws.send_str("invalid json")
|
||||
await asyncio.sleep(0.1) # Let error process
|
||||
|
||||
# Connection should remain open
|
||||
self.assertEqual(len(self.web_server.clients), 1)
|
||||
179
test/web_socket_manager_test.py
Normal file
179
test/web_socket_manager_test.py
Normal file
@@ -0,0 +1,179 @@
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
|
||||
from aiohttp.web import AppRunner
|
||||
from unittest.mock import Mock, patch, call
|
||||
import asyncio
|
||||
import unittest
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
from sia.web_agent import WebAgent, WebAgentState
|
||||
from sia.web_io_buffer import WebIOBuffer
|
||||
from sia.web_socket_manager import WebSocketManager, ClientMessage, ServerMessage
|
||||
from sia.working_memory import WorkingMemory
|
||||
from sia.system_metrics import SystemMetrics
|
||||
from sia.llm_engine import LlmEngine
|
||||
from sia.xml_validator import XMLValidator
|
||||
from sia.response_parser import ResponseParser
|
||||
|
||||
class WebSocketManagerTest(AioHTTPTestCase):
|
||||
async def get_application(self):
|
||||
"""Create test application with WebSocket handler"""
|
||||
# Create real components where beneficial for testing
|
||||
self.io_buffer = WebIOBuffer()
|
||||
self.working_memory = WorkingMemory()
|
||||
|
||||
# Create mock metrics with proper generate_context implementation
|
||||
self.mock_metrics = Mock(spec=SystemMetrics)
|
||||
metrics_elem = ET.Element("context")
|
||||
metrics_elem.attrib = {
|
||||
"time": "2024-10-31T12:00:00Z",
|
||||
"cpu": "10",
|
||||
"gpu": "20",
|
||||
"memory_used": "1000",
|
||||
"memory_total": "2000",
|
||||
"disk_used": "5000",
|
||||
"disk_total": "10000",
|
||||
"context": "0",
|
||||
"stdin": "0"
|
||||
}
|
||||
self.mock_metrics.generate_context.return_value = metrics_elem
|
||||
|
||||
# Create minimal mocks for other components
|
||||
self.mock_llm = Mock(spec=LlmEngine)
|
||||
def mock_infer(prompt: str, context: str):
|
||||
yield "<reasoning>test reasoning</reasoning>"
|
||||
self.mock_llm.infer.side_effect = mock_infer
|
||||
self.mock_validator = Mock(spec=XMLValidator)
|
||||
self.mock_validator.validate.return_value = None
|
||||
|
||||
# Create parser with real IO buffer
|
||||
self.parser = ResponseParser(self.io_buffer)
|
||||
|
||||
# Create agent with mix of real/mock components
|
||||
self.agent = WebAgent(
|
||||
system_prompt="test prompt",
|
||||
action_schema="test schema",
|
||||
working_memory=self.working_memory,
|
||||
system_metrics=self.mock_metrics,
|
||||
llm=self.mock_llm,
|
||||
validator=self.mock_validator,
|
||||
parser=self.parser
|
||||
)
|
||||
|
||||
# Create WebSocketManager
|
||||
self.ws_manager = WebSocketManager(self.agent, self.io_buffer)
|
||||
|
||||
# Create application
|
||||
app = web.Application()
|
||||
app.router.add_get('/ws', self.ws_manager.handle_websocket)
|
||||
return app
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_initial_connection(self):
|
||||
"""Test initial WebSocket connection and state messages"""
|
||||
async with self.client.ws_connect('/ws') as ws:
|
||||
# Verify initial state message
|
||||
msg = await ws.receive_json()
|
||||
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(msg['state'], WebAgentState.CONTEXT_APPROVAL.name)
|
||||
|
||||
# Verify initial context message
|
||||
msg = await ws.receive_json()
|
||||
self.assertEqual(msg['type'], ServerMessage.CONTEXT_UPDATE)
|
||||
self.assertIsNotNone(msg['context'])
|
||||
|
||||
# Verify initial response message
|
||||
msg = await ws.receive_json()
|
||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||
self.assertEqual(msg['response'], '')
|
||||
self.assertIsNone(msg['validation_error'])
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_approve_context(self):
|
||||
"""Test context approval flow"""
|
||||
async with self.client.ws_connect('/ws') as ws:
|
||||
for _ in range(3):
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
await ws.send_json({
|
||||
'type': ClientMessage.APPROVE_CONTEXT
|
||||
})
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||
self.assertEqual(msg['response'], "")
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||
self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name)
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_send_input(self):
|
||||
"""Test sending input through WebSocket"""
|
||||
test_input = "test input"
|
||||
|
||||
async with self.client.ws_connect('/ws') as ws:
|
||||
for _ in range(3):
|
||||
await ws.receive_json()
|
||||
|
||||
# Send input
|
||||
await ws.send_json({
|
||||
'type': ClientMessage.SEND_INPUT,
|
||||
'input': test_input
|
||||
})
|
||||
|
||||
# Delay for processing
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Verify input was added to buffer
|
||||
self.assertEqual(self.io_buffer.read(), test_input)
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_modify_response(self):
|
||||
"""Test modifying response in RESPONSE_APPROVAL state"""
|
||||
# Set agent to RESPONSE_APPROVAL state
|
||||
async with self.client.ws_connect('/ws') as ws:
|
||||
for _ in range(3):
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
await ws.send_json({
|
||||
'type': ClientMessage.APPROVE_CONTEXT
|
||||
})
|
||||
for _ in range(4):
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
modified_response = "<reasoning>modified test</reasoning>"
|
||||
await ws.send_json({
|
||||
'type': ClientMessage.MODIFY_RESPONSE,
|
||||
'response': modified_response
|
||||
})
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||
self.assertEqual(msg['response'], modified_response)
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_multiple_clients(self):
|
||||
"""Test broadcasting to multiple WebSocket clients"""
|
||||
# Connect first client
|
||||
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
|
||||
for _ in range(3):
|
||||
await asyncio.wait_for(ws1.receive_json(), timeout=0.5)
|
||||
await asyncio.wait_for(ws2.receive_json(), timeout=0.5)
|
||||
await ws1.send_json({
|
||||
'type': ClientMessage.APPROVE_CONTEXT
|
||||
})
|
||||
msg1 = await ws1.receive_json()
|
||||
msg2 = await ws2.receive_json()
|
||||
self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(msg1['state'], WebAgentState.INFERENCE.name)
|
||||
self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up resources"""
|
||||
super().tearDown()
|
||||
self.io_buffer = None
|
||||
self.working_memory = None
|
||||
self.agent = None
|
||||
self.ws_manager = None
|
||||
Reference in New Issue
Block a user