Update unit tests
This commit is contained in:
Niels Geens
@@ -1,87 +1,114 @@
|
||||
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
|
||||
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, LlmState
|
||||
|
||||
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": LlmState.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"] = LlmState.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"] = LlmState.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"] = LlmState.INFERENCE
|
||||
self.approver._handle_llm_state_change("default", LlmState.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"] = LlmState.IDLE
|
||||
self.approver.context_enabled = True
|
||||
|
||||
# Change state quickly
|
||||
time.sleep(0.1)
|
||||
self.mock_agent.llms["default"] = LlmState.INFERENCE
|
||||
self.approver._handle_llm_state_change("default", LlmState.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"] = LlmState.IDLE
|
||||
|
||||
# Create a new instance with a fresh mock for this test
|
||||
mock_agent = Mock(spec=WebAgent)
|
||||
mock_agent.llms = {"default": LlmState.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"] = LlmState.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
|
||||
Reference in New Issue
Block a user