246 lines
9.6 KiB
Python
246 lines
9.6 KiB
Python
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 get_metrics implementation
|
|
self.mock_metrics = Mock(spec=SystemMetrics)
|
|
self.mock_metrics.get_metrics.return_value = {
|
|
"timestamp": "2024-10-31T12:00:00Z",
|
|
"cpu": 10,
|
|
"gpu": 20,
|
|
"memory_used": 1000,
|
|
"memory_total": 2000,
|
|
"disk_used": 5000,
|
|
"disk_total": 10000,
|
|
"token_count": 100,
|
|
"token_limit": 1000
|
|
}
|
|
|
|
# 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
|
|
self.mock_llm.token_count.return_value = 100
|
|
self.mock_llm.token_limit.return_value = 1000
|
|
|
|
# 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,
|
|
metrics=self.mock_metrics,
|
|
llm=self.mock_llm,
|
|
validator=self.mock_validator,
|
|
parser=self.parser
|
|
)
|
|
|
|
# Create WebSocketManager
|
|
self.ws_manager = await WebSocketManager.create(self.agent, self.io_buffer)
|
|
|
|
# Create application
|
|
app = web.Application()
|
|
app.router.add_get('/ws', self.ws_manager.handle_websocket)
|
|
return app
|
|
|
|
async def _skip_initial_messages(self, ws):
|
|
"""Helper to skip past the initial connection messages"""
|
|
for _ in range(4): # Now 4 messages including auto approver config
|
|
await ws.receive_json()
|
|
|
|
@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'])
|
|
|
|
# Verify initial auto approver config
|
|
msg = await ws.receive_json()
|
|
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
|
self.assertIn('config', msg)
|
|
|
|
@unittest_run_loop
|
|
async def test_approve_context(self):
|
|
"""Test context approval flow"""
|
|
async with self.client.ws_connect('/ws') as ws:
|
|
await self._skip_initial_messages(ws)
|
|
|
|
await ws.send_json({
|
|
'type': ClientMessage.APPROVE_CONTEXT
|
|
})
|
|
|
|
msg = await ws.receive_json()
|
|
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
|
self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
|
|
|
|
msg = await ws.receive_json()
|
|
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
|
self.assertEqual(msg['response'], "")
|
|
|
|
msg = await ws.receive_json()
|
|
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
|
self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
|
|
|
|
msg = await ws.receive_json()
|
|
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:
|
|
await self._skip_initial_messages(ws)
|
|
|
|
await ws.send_json({
|
|
'type': ClientMessage.SEND_INPUT,
|
|
'input': test_input
|
|
})
|
|
|
|
await asyncio.sleep(0.5)
|
|
self.assertEqual(self.io_buffer.read(), test_input)
|
|
|
|
@unittest_run_loop
|
|
async def test_modify_response(self):
|
|
"""Test modifying response in RESPONSE_APPROVAL state"""
|
|
async with self.client.ws_connect('/ws') as ws:
|
|
await self._skip_initial_messages(ws)
|
|
|
|
await ws.send_json({
|
|
'type': ClientMessage.APPROVE_CONTEXT
|
|
})
|
|
|
|
# Skip to RESPONSE_APPROVAL state
|
|
for _ in range(4):
|
|
await ws.receive_json()
|
|
|
|
modified_response = "<reasoning>modified test</reasoning>"
|
|
await ws.send_json({
|
|
'type': ClientMessage.MODIFY_RESPONSE,
|
|
'response': modified_response
|
|
})
|
|
|
|
msg = await ws.receive_json()
|
|
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"""
|
|
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
|
|
await self._skip_initial_messages(ws1)
|
|
await self._skip_initial_messages(ws2)
|
|
|
|
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)
|
|
|
|
@unittest_run_loop
|
|
async def test_auto_approver_config(self):
|
|
"""Test setting and receiving auto approver config"""
|
|
async with self.client.ws_connect('/ws') as ws:
|
|
await self._skip_initial_messages(ws)
|
|
|
|
test_config = {
|
|
'context_enabled': True,
|
|
'response_enabled': True,
|
|
'context_timeout': 10.0,
|
|
'response_timeout': 20.0
|
|
}
|
|
|
|
await ws.send_json({
|
|
'type': ClientMessage.AUTO_APPROVER_CONFIG,
|
|
'config': test_config
|
|
})
|
|
|
|
msg = await ws.receive_json()
|
|
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
|
self.assertEqual(msg['config'], test_config)
|
|
|
|
self.ws_manager.auto_approver.context_enabled = False
|
|
self.ws_manager.auto_approver.response_enabled = False
|
|
|
|
@unittest_run_loop
|
|
async def test_auto_approver_broadcast(self):
|
|
"""Test broadcasting auto approver config changes to multiple clients"""
|
|
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
|
|
await self._skip_initial_messages(ws1)
|
|
await self._skip_initial_messages(ws2)
|
|
|
|
test_config = {
|
|
'context_enabled': True,
|
|
'response_enabled': False,
|
|
'context_timeout': 15.0,
|
|
'response_timeout': 25.0
|
|
}
|
|
|
|
await ws1.send_json({
|
|
'type': ClientMessage.AUTO_APPROVER_CONFIG,
|
|
'config': test_config
|
|
})
|
|
|
|
msg1 = await ws1.receive_json()
|
|
msg2 = await ws2.receive_json()
|
|
self.assertEqual(msg1['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
|
self.assertEqual(msg2['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
|
self.assertEqual(msg1['config'], test_config)
|
|
self.assertEqual(msg2['config'], test_config)
|
|
|
|
self.ws_manager.auto_approver.context_enabled = False
|
|
self.ws_manager.auto_approver.response_enabled = False
|
|
|
|
def tearDown(self):
|
|
"""Clean up resources"""
|
|
super().tearDown()
|
|
self.io_buffer = None
|
|
self.working_memory = None
|
|
self.agent = None
|
|
self.ws_manager = None
|