Add auto_approver
This commit is contained in:
Niels Geens
87
test/auto_approver_test.py
Normal file
87
test/auto_approver_test.py
Normal file
@@ -0,0 +1,87 @@
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
from threading import Event
|
||||
import time
|
||||
|
||||
from sia.auto_approver import AutoApprover
|
||||
from sia.web_agent import WebAgent, WebAgentState
|
||||
|
||||
class AutoApproverTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Create mock agent and approver for each test"""
|
||||
self.mock_agent = Mock(spec=WebAgent)
|
||||
self.mock_agent.state = WebAgentState.UPDATE
|
||||
self.approver = AutoApprover(self.mock_agent)
|
||||
|
||||
def test_timeout_setters(self):
|
||||
"""Test timeout setters block changes when enabled"""
|
||||
self.approver.context_timeout = 10.0
|
||||
self.approver.response_timeout = 15.0
|
||||
|
||||
self.approver.context_enabled = True
|
||||
with self.assertRaises(ValueError):
|
||||
self.approver.context_timeout = 20.0
|
||||
|
||||
self.approver.response_enabled = True
|
||||
with self.assertRaises(ValueError):
|
||||
self.approver.response_timeout = 25.0
|
||||
|
||||
def test_enable_starts_thread_in_correct_state(self):
|
||||
"""Test enabling only starts thread in matching state"""
|
||||
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||
self.approver.context_enabled = True
|
||||
time.sleep(0.1) # Allow thread to start
|
||||
self.assertTrue(self.approver._context_thread.is_alive())
|
||||
self.assertIsNone(self.approver._response_thread)
|
||||
|
||||
def test_state_change_stops_threads(self):
|
||||
"""Test state changes stop irrelevant threads"""
|
||||
# Start in context approval with thread
|
||||
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||
self.approver.context_enabled = True
|
||||
time.sleep(0.1)
|
||||
self.assertTrue(self.approver._context_thread.is_alive())
|
||||
|
||||
# Change state and verify thread stops
|
||||
self.mock_agent.state = WebAgentState.INFERENCE
|
||||
self.approver._handle_state_change(WebAgentState.INFERENCE)
|
||||
time.sleep(0.1)
|
||||
self.assertIsNone(self.approver._context_thread)
|
||||
|
||||
def test_quick_state_cycle(self):
|
||||
"""Test thread stops when state changes before timeout"""
|
||||
self.approver.context_timeout = 5.0
|
||||
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||
self.approver.context_enabled = True
|
||||
|
||||
# Change state quickly
|
||||
time.sleep(0.1)
|
||||
self.mock_agent.state = WebAgentState.INFERENCE
|
||||
self.approver._handle_state_change(WebAgentState.INFERENCE)
|
||||
|
||||
# Verify no approval happened
|
||||
self.mock_agent.approve_context.assert_not_called()
|
||||
|
||||
def test_disable_stops_thread(self):
|
||||
"""Test disabling stops active thread"""
|
||||
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||
self.approver.context_enabled = True
|
||||
time.sleep(0.1)
|
||||
self.assertTrue(self.approver._context_thread.is_alive())
|
||||
|
||||
self.approver.context_enabled = False
|
||||
self.assertIsNone(self.approver._context_thread)
|
||||
|
||||
def test_successful_auto_approval(self):
|
||||
"""Test thread approves after timeout when conditions met"""
|
||||
self.approver.context_timeout = 0.1
|
||||
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||
self.approver.context_enabled = True
|
||||
|
||||
time.sleep(0.2)
|
||||
self.mock_agent.approve_context.assert_called_once()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up any running threads"""
|
||||
self.approver.context_enabled = False
|
||||
self.approver.response_enabled = False
|
||||
@@ -68,6 +68,11 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
||||
app.router.add_get('/ws', self.ws_manager.handle_websocket)
|
||||
return app
|
||||
|
||||
async def _skip_initial_messages(self, ws):
|
||||
"""Helper to skip past the initial connection messages"""
|
||||
for _ in range(4): # Now 4 messages including auto approver config
|
||||
await ws.receive_json()
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_initial_connection(self):
|
||||
"""Test initial WebSocket connection and state messages"""
|
||||
@@ -88,25 +93,34 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
||||
self.assertEqual(msg['response'], '')
|
||||
self.assertIsNone(msg['validation_error'])
|
||||
|
||||
# Verify initial auto approver config
|
||||
msg = await ws.receive_json()
|
||||
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||
self.assertIn('config', msg)
|
||||
|
||||
@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 self._skip_initial_messages(ws)
|
||||
|
||||
await ws.send_json({
|
||||
'type': ClientMessage.APPROVE_CONTEXT
|
||||
})
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
|
||||
msg = await ws.receive_json()
|
||||
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)
|
||||
|
||||
msg = await ws.receive_json()
|
||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||
self.assertEqual(msg['response'], "")
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
|
||||
msg = await ws.receive_json()
|
||||
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)
|
||||
|
||||
msg = await ws.receive_json()
|
||||
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name)
|
||||
|
||||
@@ -116,53 +130,51 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
||||
test_input = "test input"
|
||||
|
||||
async with self.client.ws_connect('/ws') as ws:
|
||||
for _ in range(3):
|
||||
await ws.receive_json()
|
||||
await self._skip_initial_messages(ws)
|
||||
|
||||
# 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 self._skip_initial_messages(ws)
|
||||
|
||||
await ws.send_json({
|
||||
'type': ClientMessage.APPROVE_CONTEXT
|
||||
})
|
||||
|
||||
# Skip to RESPONSE_APPROVAL state
|
||||
for _ in range(4):
|
||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
||||
await ws.receive_json()
|
||||
|
||||
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)
|
||||
|
||||
msg = await ws.receive_json()
|
||||
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 self._skip_initial_messages(ws1)
|
||||
await self._skip_initial_messages(ws2)
|
||||
|
||||
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)
|
||||
@@ -170,6 +182,60 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
||||
self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE)
|
||||
self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name)
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_auto_approver_config(self):
|
||||
"""Test setting and receiving auto approver config"""
|
||||
async with self.client.ws_connect('/ws') as ws:
|
||||
await self._skip_initial_messages(ws)
|
||||
|
||||
test_config = {
|
||||
'context_enabled': True,
|
||||
'response_enabled': True,
|
||||
'context_timeout': 10.0,
|
||||
'response_timeout': 20.0
|
||||
}
|
||||
|
||||
await ws.send_json({
|
||||
'type': ClientMessage.AUTO_APPROVER_CONFIG,
|
||||
'config': test_config
|
||||
})
|
||||
|
||||
msg = await ws.receive_json()
|
||||
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||
self.assertEqual(msg['config'], test_config)
|
||||
|
||||
self.ws_manager.auto_approver.context_enabled = False
|
||||
self.ws_manager.auto_approver.response_enabled = False
|
||||
|
||||
@unittest_run_loop
|
||||
async def test_auto_approver_broadcast(self):
|
||||
"""Test broadcasting auto approver config changes to multiple clients"""
|
||||
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
|
||||
await self._skip_initial_messages(ws1)
|
||||
await self._skip_initial_messages(ws2)
|
||||
|
||||
test_config = {
|
||||
'context_enabled': True,
|
||||
'response_enabled': False,
|
||||
'context_timeout': 15.0,
|
||||
'response_timeout': 25.0
|
||||
}
|
||||
|
||||
await ws1.send_json({
|
||||
'type': ClientMessage.AUTO_APPROVER_CONFIG,
|
||||
'config': test_config
|
||||
})
|
||||
|
||||
msg1 = await ws1.receive_json()
|
||||
msg2 = await ws2.receive_json()
|
||||
self.assertEqual(msg1['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||
self.assertEqual(msg2['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||
self.assertEqual(msg1['config'], test_config)
|
||||
self.assertEqual(msg2['config'], test_config)
|
||||
|
||||
self.ws_manager.auto_approver.context_enabled = False
|
||||
self.ws_manager.auto_approver.response_enabled = False
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up resources"""
|
||||
super().tearDown()
|
||||
|
||||
Reference in New Issue
Block a user