252 lines
9.5 KiB
Python
252 lines
9.5 KiB
Python
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)
|