Approve context

test/web_agent_test.py

@@ -1,3 +1,4 @@
+import asyncio
 import unittest
 from datetime import datetime
 from unittest.mock import Mock, MagicMock, patch, call
@@ -123,7 +124,7 @@ class WebAgentTest(unittest.TestCase):
         self.agent._current_state = WebAgentState.UPDATE
         
         with self.assertRaises(Exception) as context:
-            self.agent.approve_context()
+            asyncio.run(self.agent.approve_context())
         self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception))
         
     def test_approve_response_state_error(self):
@@ -132,7 +133,7 @@ class WebAgentTest(unittest.TestCase):
             self.agent.approve_response()
         self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
         
-    def test_context_approval_flow(self):
+    async def test_context_approval_flow(self):
         """Test complete context approval flow with state transitions."""
         # Register handlers
         self.agent.add_state_change_handler(self.state_change_handler)
@@ -140,6 +141,7 @@ class WebAgentTest(unittest.TestCase):
         
         # Setup LLM response
         def mock_infer(prompt: str, context: str) -> Iterator[str]:
+            # Return complete response at once for testing
             yield "<reasoning>test reasoning</reasoning>"
         self.mock_llm.infer.side_effect = mock_infer
         
@@ -149,6 +151,9 @@ class WebAgentTest(unittest.TestCase):
         # Approve context
         self.agent.approve_context()
         
+        # Wait for response tasks to complete
+        await asyncio.gather(*self.web_server._response_tasks)
+        
         # Verify state transitions
         expected_states = [
             WebAgentState.INFERENCE,
@@ -156,10 +161,11 @@ class WebAgentTest(unittest.TestCase):
         ]
         self.assertEqual(self.state_changes, expected_states)
         
-        # Verify response updates
-        self.assertEqual(len(self.response_changes), 2)  # Empty + final response
-        self.assertEqual(self.response_changes[0], "")
-        self.assertEqual(self.response_changes[1], "<reasoning>test reasoning</reasoning>")
+        # Updated test: Should now expect streaming updates
+        self.assertEqual(len(self.response_changes), len(self.received_messages))
+        for i, msg in enumerate(self.received_messages):
+            self.assertEqual(msg["type"], "RESPONSE_UPDATE")
+            self.assertEqual(msg["data"], self.response_changes[i])
         
     def test_response_approval_flow_command(self):
         """Test response approval flow with command execution."""
@@ -200,7 +206,7 @@ class WebAgentTest(unittest.TestCase):
         
         # Set initial state and response
         self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
-        self.agent._response = "<reasoning>test reasoning</reasoning>"
+        self.agent.response = "<reasoning>test reasoning</reasoning>"
         
         # Approve response
         self.agent.approve_response()
@@ -218,7 +224,7 @@ class WebAgentTest(unittest.TestCase):
         ]
         self.assertEqual(self.state_changes, expected_states)
         
-    def test_response_validation(self):
+    async def test_response_validation(self):
         """Test response validation during response setting."""
         # Register handlers
         self.agent.add_response_change_handler(self.response_change_handler)
@@ -227,8 +233,8 @@ class WebAgentTest(unittest.TestCase):
         error_message = "Invalid XML"
         self.mock_validator.validate.return_value = error_message
         
-        # Set response
-        self.agent._set_response("<invalid>")
+        # Set response using async method
+        await self.agent._set_response("<invalid>")
         
         # Verify validation error stored
         self.assertEqual(self.agent._validation_error, error_message)