Separate and tested web_socket_manager

test/web_agent_test.py

@@ -88,7 +88,7 @@ class WebAgentTest(unittest.TestCase):
         """Track state changes for verification."""
         self.state_changes.append(state)
         
-    def response_change_handler(self, response: str):
+    def response_change_handler(self, response: str, validation_error: str):
         """Track response changes for verification."""
         self.response_changes.append(response)
         
@@ -97,7 +97,7 @@ class WebAgentTest(unittest.TestCase):
         self.assertEqual(self.agent.current_state, WebAgentState.CONTEXT_APPROVAL)
         self.assertEqual(self.agent.response, "")
         self.assertIsNotNone(self.agent.context)
-        self.assertIsNone(self.agent._validation_error)
+        self.assertIsNone(self.agent.validation_error)
         self.assertEqual(len(self.agent._state_change_handlers), 0)
         self.assertEqual(len(self.agent._response_change_handlers), 0)
         
@@ -133,7 +133,7 @@ class WebAgentTest(unittest.TestCase):
             self.agent.approve_response()
         self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
         
-    async def test_context_approval_flow(self):
+    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)
@@ -145,27 +145,16 @@ class WebAgentTest(unittest.TestCase):
             yield "<reasoning>test reasoning</reasoning>"
         self.mock_llm.infer.side_effect = mock_infer
         
-        # Set initial state
-        self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
-        
         # 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,
             WebAgentState.RESPONSE_APPROVAL
         ]
         self.assertEqual(self.state_changes, expected_states)
-        
-        # 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])
+        self.assertEqual(self.response_changes[-1], "<reasoning>test reasoning</reasoning>")
         
     def test_response_approval_flow_command(self):
         """Test response approval flow with command execution."""
@@ -206,7 +195,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.set_response("<reasoning>test reasoning</reasoning>")
         
         # Approve response
         self.agent.approve_response()
@@ -224,22 +213,13 @@ class WebAgentTest(unittest.TestCase):
         ]
         self.assertEqual(self.state_changes, expected_states)
         
-    async def test_response_validation(self):
+    def test_response_validation(self):
         """Test response validation during response setting."""
-        # Register handlers
         self.agent.add_response_change_handler(self.response_change_handler)
-        
-        # Mock validator to fail
         error_message = "Invalid XML"
         self.mock_validator.validate.return_value = error_message
-        
-        # Set response using async method
-        await self.agent._set_response("<invalid>")
-        
-        # Verify validation error stored
-        self.assertEqual(self.agent._validation_error, error_message)
-        
-        # Verify response updated
+        self.agent.set_response("<invalid>")
+        self.assertEqual(self.agent.validation_error, error_message)
         self.assertEqual(self.response_changes, ["<invalid>"])
         
     def tearDown(self):