114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
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, AgentState
|
|
|
|
class AutoApproverTest(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Create mock agent and approver for each test"""
|
|
self.mock_agent = Mock(spec=WebAgent)
|
|
# Set up llms attribute as a dictionary
|
|
self.mock_agent.llms = {"default": AgentState.IDLE}
|
|
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"""
|
|
# Mock the agent's llm state
|
|
self.mock_agent.llms["default"] = AgentState.IDLE
|
|
self.approver.context_enabled = True
|
|
time.sleep(0.1) # Allow thread to start
|
|
self.assertIsNotNone(self.approver._context_thread)
|
|
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 idle state with context thread
|
|
self.mock_agent.llms["default"] = AgentState.IDLE
|
|
|
|
# Mock _stop_context_thread to verify it's called
|
|
with patch.object(self.approver, '_stop_context_thread') as mock_stop:
|
|
self.approver.context_enabled = True
|
|
time.sleep(0.1)
|
|
|
|
# Change state and verify stop method was called
|
|
self.mock_agent.llms["default"] = AgentState.INFERENCE
|
|
self.approver._handle_state_change(AgentState.INFERENCE)
|
|
mock_stop.assert_called_once()
|
|
|
|
def test_quick_state_cycle(self):
|
|
"""Test thread stops when state changes before timeout"""
|
|
self.approver.context_timeout = 5.0
|
|
self.mock_agent.llms["default"] = AgentState.IDLE
|
|
self.approver.context_enabled = True
|
|
|
|
# Change state quickly
|
|
time.sleep(0.1)
|
|
self.mock_agent.llms["default"] = AgentState.INFERENCE
|
|
self.approver._handle_state_change(AgentState.INFERENCE)
|
|
|
|
# Verify no approval happened
|
|
self.mock_agent.run_inference.assert_not_called()
|
|
|
|
def test_disable_stops_thread(self):
|
|
"""Test disabling stops active thread"""
|
|
self.mock_agent.llms["default"] = AgentState.IDLE
|
|
|
|
# Create a new instance with a fresh mock for this test
|
|
mock_agent = Mock(spec=WebAgent)
|
|
mock_agent.llms = {"default": AgentState.IDLE}
|
|
test_approver = AutoApprover(mock_agent)
|
|
|
|
# Start with a mocked stop method
|
|
with patch.object(test_approver, '_stop_context_thread') as mock_stop:
|
|
# Enable context
|
|
test_approver.context_enabled = True
|
|
time.sleep(0.1)
|
|
|
|
# Reset the mock to clear previous calls
|
|
mock_stop.reset_mock()
|
|
|
|
# Disable and verify stop method was called exactly once
|
|
test_approver.context_enabled = False
|
|
mock_stop.assert_called_once()
|
|
|
|
def test_successful_auto_approval(self):
|
|
"""Test thread approves after timeout when conditions met"""
|
|
self.approver.context_timeout = 0.1
|
|
self.mock_agent.llms["default"] = AgentState.IDLE
|
|
self.approver.llm_name = "default"
|
|
|
|
# Use threading Event to properly synchronize the test
|
|
approval_event = Event()
|
|
|
|
def mock_run_inference(llm_name):
|
|
approval_event.set()
|
|
|
|
self.mock_agent.run_inference = mock_run_inference
|
|
|
|
self.approver.context_enabled = True
|
|
|
|
# Wait for approval with timeout
|
|
approved = approval_event.wait(0.5)
|
|
self.assertTrue(approved, "Auto-approval did not occur within expected time")
|
|
|
|
def tearDown(self):
|
|
"""Clean up any running threads"""
|
|
self.approver.context_enabled = False
|
|
self.approver.response_enabled = False |