Separate and tested web_socket_manager

test/web_socket_manager_test.py

@@ -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