Approve context

test/web_server_test.py

@@ -6,6 +6,7 @@ from sia.web_agent import WebAgentState
 from unittest.mock import patch
 import asyncio
 import asyncio
+import time
 import json
 import json
 import shutil
@@ -29,6 +30,7 @@ class WebServerTest(AioHTTPTestCase):
         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,
@@ -54,7 +56,6 @@ class WebServerTest(AioHTTPTestCase):
                 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}")
@@ -108,6 +109,7 @@ class WebServerTest(AioHTTPTestCase):
             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)
@@ -124,7 +126,7 @@ class WebServerTest(AioHTTPTestCase):
             self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE)
             self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name)
 
-    @unittest_run_loop        
+    @unittest_run_loop
     async def test_validation_error(self):
         """Test handling of validation errors."""
         async with self.client.ws_connect("/ws") as ws:
@@ -134,7 +136,7 @@ class WebServerTest(AioHTTPTestCase):
             
             # Set validation error and response
             self.agent._validation_error = "Invalid XML"
-            await self.agent._set_response("<invalid>")
+            self.agent._set_response("<invalid>")  # Now synchronous
             await asyncio.sleep(0.1)  # Let changes propagate
             
             # Should receive response update with error
@@ -204,24 +206,34 @@ class WebServerTest(AioHTTPTestCase):
             await ws1.close()
             await ws2.close()
             
-    @unittest_run_loop
+    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):
-        """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
+            await asyncio.sleep(0.1)
+
             test_input = "test input"
             await ws.send_json({
                 "type": ClientMessage.SEND_INPUT,
                 "data": test_input
             })
-            await asyncio.sleep(0.1)  # Let input process
+
+            # Wait for buffer to have content
+            self.assertTrue(await self.wait_for_buffer_content())
             
-            # Verify input added to buffer with newline
-            self.assertEqual(self.io_buffer.read(), test_input + "\n")
+            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):