Implemented webserver
This commit is contained in:
228
test/web_server_test.py
Normal file
228
test/web_server_test.py
Normal file
@@ -0,0 +1,228 @@
|
||||
from aiohttp import WSMsgType
|
||||
from aiohttp import WSMsgType
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
|
||||
from sia.web_agent import WebAgentState
|
||||
from typing import AsyncIterator, Iterator
|
||||
from unittest.mock import Mock, patch
|
||||
import asyncio
|
||||
import asyncio
|
||||
import json
|
||||
import json
|
||||
|
||||
from sia.llm_engine import LlmEngine
|
||||
from sia.response_parser import ResponseParser
|
||||
from sia.system_metrics import SystemMetrics
|
||||
from sia.web_agent import WebAgentState
|
||||
from sia.web_io_buffer import WebIOBuffer
|
||||
from sia.web_server import WebServer, ServerMessage, ClientMessage
|
||||
from sia.working_memory import WorkingMemory
|
||||
from sia.xml_validator import XMLValidator
|
||||
from .mock_web_agent import MockWebAgent
|
||||
|
||||
class WebServerTest(AioHTTPTestCase):
|
||||
async def get_application(self):
|
||||
"""Create application for testing."""
|
||||
self.io_buffer = WebIOBuffer()
|
||||
|
||||
# Create MockWebAgent instance
|
||||
self.agent = MockWebAgent()
|
||||
|
||||
# Create web server
|
||||
self.web_server = WebServer(self.agent, self.io_buffer, "localhost", 8080)
|
||||
return self.web_server.app
|
||||
|
||||
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:
|
||||
print("WebSocket closed normally")
|
||||
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,
|
||||
})
|
||||
|
||||
# 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"
|
||||
await self.agent._set_response("<invalid>")
|
||||
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()
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_send_input(self):
|
||||
"""Test sending input to agent."""
|
||||
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 input
|
||||
test_input = "test input"
|
||||
await ws.send_json({
|
||||
"type": ClientMessage.SEND_INPUT,
|
||||
"data": test_input
|
||||
})
|
||||
await asyncio.sleep(0.1) # Let input process
|
||||
|
||||
# Verify input added to buffer with newline
|
||||
self.assertEqual(self.io_buffer.read(), test_input + "\n")
|
||||
|
||||
@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)
|
||||
Reference in New Issue
Block a user