Update unit tests
This commit is contained in:
Niels Geens
@@ -1,178 +1,170 @@
|
||||
import unittest
|
||||
from datetime import datetime
|
||||
from unittest.mock import Mock, MagicMock, patch
|
||||
import xml.etree.ElementTree as ET
|
||||
import logging
|
||||
|
||||
from sia.web_io_buffer import WebIOBuffer
|
||||
from sia.working_memory import WorkingMemory
|
||||
from sia.system_metrics import SystemMetrics
|
||||
from sia.llm_engine import LlmEngine
|
||||
from sia.xml_validator import XMLValidator
|
||||
from sia.response_parser import ResponseParser
|
||||
from sia.base_agent import BaseAgent
|
||||
from sia.command import Command
|
||||
from sia.command_result import CommandResult
|
||||
from sia.reasoning_entry import ReasoningEntry
|
||||
from sia.entry.parse_error_entry import ParseErrorEntry
|
||||
from sia.delete_command import DeleteCommand
|
||||
from sia.entry import Entry
|
||||
|
||||
class TestLogHandler(logging.Handler):
|
||||
"""Custom logging handler that stores records for test verification."""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.records = []
|
||||
|
||||
def emit(self, record):
|
||||
self.records.append(record)
|
||||
|
||||
def clear(self):
|
||||
self.records = []
|
||||
|
||||
class MockEntry(Entry):
|
||||
"""Mock entry class that properly extends Entry."""
|
||||
def __init__(self, id: str, timestamp: datetime, update_behavior=None):
|
||||
super().__init__(id, timestamp)
|
||||
self._update_behavior = update_behavior
|
||||
self.update_called = False
|
||||
|
||||
def update(self) -> None:
|
||||
self.update_called = True
|
||||
if self._update_behavior:
|
||||
self._update_behavior()
|
||||
|
||||
def generate_context(self) -> ET.Element:
|
||||
elem = ET.Element("mock_entry", {"id": self.id})
|
||||
elem.text = "<![CDATA[mock content]]>"
|
||||
return elem
|
||||
|
||||
class TestBaseAgent(BaseAgent):
|
||||
"""Concrete implementation of BaseAgent for testing."""
|
||||
pass
|
||||
|
||||
class BaseAgentTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test cases with mocked components."""
|
||||
self.mock_llm = Mock(spec=LlmEngine)
|
||||
self.mock_llm.token_count.return_value = 100
|
||||
self.mock_llm.token_limit.return_value = 1000
|
||||
self.mock_validator = Mock(spec=XMLValidator)
|
||||
self.io_buffer = WebIOBuffer()
|
||||
self.working_memory = WorkingMemory()
|
||||
self.parser = ResponseParser(self.io_buffer)
|
||||
|
||||
# Mock system metrics
|
||||
self.mock_metrics = Mock(spec=SystemMetrics)
|
||||
self.mock_metrics.get_metrics.return_value = {
|
||||
"timestamp": "2024-10-31T12:00:00Z",
|
||||
"cpu": 10,
|
||||
"gpu": 20,
|
||||
"memory_used": 1000,
|
||||
"memory_total": 2000,
|
||||
"disk_used": 5000,
|
||||
"disk_total": 10000
|
||||
}
|
||||
|
||||
# Create test agent
|
||||
self.agent = TestBaseAgent(
|
||||
system_prompt="test prompt",
|
||||
action_schema="test schema",
|
||||
working_memory=self.working_memory,
|
||||
metrics=self.mock_metrics,
|
||||
llm=self.mock_llm,
|
||||
validator=self.mock_validator,
|
||||
parser=self.parser
|
||||
)
|
||||
|
||||
# Set timestamp for entries
|
||||
self.test_timestamp = datetime(2024, 10, 31, 12, 0, 0)
|
||||
|
||||
# Setup logging
|
||||
self.log_handler = TestLogHandler()
|
||||
logger = logging.getLogger()
|
||||
logger.addHandler(self.log_handler)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up resources."""
|
||||
logger = logging.getLogger()
|
||||
logger.removeHandler(self.log_handler)
|
||||
|
||||
if hasattr(self, 'agent'):
|
||||
del self.agent
|
||||
|
||||
def test_initialization(self):
|
||||
"""Test agent initialization and component setup."""
|
||||
self.assertEqual(self.agent._working_memory, self.working_memory)
|
||||
self.assertEqual(self.agent._metrics, self.mock_metrics)
|
||||
self.assertEqual(self.agent._llm, self.mock_llm)
|
||||
self.assertEqual(self.agent._validator, self.mock_validator)
|
||||
self.assertEqual(self.agent._parser, self.parser)
|
||||
self.assertEqual(self.agent._action_schema, "test schema")
|
||||
|
||||
def test_cleanup(self):
|
||||
"""Test cleanup on agent deletion."""
|
||||
del self.agent
|
||||
self.mock_metrics.stop.assert_called_once()
|
||||
|
||||
def test_compile_context_empty_memory(self):
|
||||
"""Test context compilation with empty working memory."""
|
||||
context = self.agent._compile_context()
|
||||
|
||||
# Parse context and verify structure
|
||||
root = ET.fromstring(context)
|
||||
self.assertEqual(root.tag, "context")
|
||||
|
||||
# Check metrics
|
||||
self.assertEqual(root.get("cpu"), "10")
|
||||
self.assertEqual(root.get("gpu"), "20")
|
||||
self.assertEqual(root.get("memory_used"), "1000")
|
||||
self.assertEqual(root.get("memory_total"), "2000")
|
||||
self.assertEqual(root.get("disk_used"), "5000")
|
||||
self.assertEqual(root.get("disk_total"), "10000")
|
||||
|
||||
# Check context size is 0 (empty memory)
|
||||
self.assertEqual(root.get("context"), "10.0")
|
||||
|
||||
# Check no memory entries
|
||||
self.assertEqual(len(list(root)), 0)
|
||||
|
||||
def test_compile_context_with_entries(self):
|
||||
"""Test context compilation with a single entry."""
|
||||
entry = ReasoningEntry("test-id", self.test_timestamp, "test reasoning")
|
||||
self.agent._working_memory.add_entry(entry)
|
||||
|
||||
context = self.agent._compile_context()
|
||||
root = ET.fromstring(context)
|
||||
|
||||
# Check context size reflects one entry
|
||||
self.assertEqual(root.get("context"), "10.0")
|
||||
|
||||
# Verify entry content
|
||||
reasoning_elem = root.find("reasoning")
|
||||
self.assertIsNotNone(reasoning_elem)
|
||||
self.assertEqual(reasoning_elem.get("id"), "test-id")
|
||||
self.assertEqual(reasoning_elem.text.strip(), "test reasoning")
|
||||
|
||||
def test_multiple_entries(self):
|
||||
"""Test handling multiple entries in working memory."""
|
||||
entries = [
|
||||
ReasoningEntry(f"id-{i}", self.test_timestamp, f"test {i}")
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.agent._working_memory.add_entry(entry)
|
||||
|
||||
context = self.agent._compile_context()
|
||||
root = ET.fromstring(context)
|
||||
|
||||
# Check context size reflects three entries
|
||||
self.assertEqual(root.get("context"), "10.0")
|
||||
|
||||
# Verify entry order maintained
|
||||
reasoning_elems = root.findall("reasoning")
|
||||
self.assertEqual(len(reasoning_elems), 3)
|
||||
for i, elem in enumerate(reasoning_elems):
|
||||
self.assertEqual(elem.get("id"), f"id-{i}")
|
||||
self.assertEqual(elem.text.strip(), f"test {i}")
|
||||
import unittest
|
||||
from unittest.mock import Mock, MagicMock, patch
|
||||
import xml.etree.ElementTree as ET
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from sia.web_io_buffer import WebIOBuffer
|
||||
from sia.working_memory import WorkingMemory
|
||||
from sia.system_metrics import SystemMetrics
|
||||
from sia.llm_engine import LlmEngine
|
||||
from sia.xml_validator import XMLValidator
|
||||
from sia.response_parser import ResponseParser
|
||||
from sia.base_agent import BaseAgent
|
||||
from sia.command import Command
|
||||
from sia.command_result import CommandResult
|
||||
from sia.entry.reasoning_entry import ReasoningEntry
|
||||
from sia.entry.parse_error_entry import ParseErrorEntry
|
||||
from sia.delete_command import DeleteCommand
|
||||
from sia.entry import Entry
|
||||
|
||||
class TestLogHandler(logging.Handler):
|
||||
"""Custom logging handler that stores records for test verification."""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.records = []
|
||||
|
||||
def emit(self, record):
|
||||
self.records.append(record)
|
||||
|
||||
def clear(self):
|
||||
self.records = []
|
||||
|
||||
class MockEntry(Entry):
|
||||
"""Mock entry class that properly extends Entry."""
|
||||
def __init__(self, id: str, update_behavior=None):
|
||||
super().__init__(id)
|
||||
self._update_behavior = update_behavior
|
||||
self.update_called = False
|
||||
|
||||
def update(self) -> None:
|
||||
self.update_called = True
|
||||
if self._update_behavior:
|
||||
self._update_behavior()
|
||||
|
||||
def generate_context(self) -> ET.Element:
|
||||
elem = ET.Element("mock_entry", {"id": self.id})
|
||||
elem.text = "<![CDATA[mock content]]>"
|
||||
return elem
|
||||
|
||||
class TestBaseAgent(BaseAgent):
|
||||
"""Concrete implementation of BaseAgent for testing."""
|
||||
pass
|
||||
|
||||
class BaseAgentTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test cases with mocked components."""
|
||||
self.mock_llm = Mock(spec=LlmEngine)
|
||||
self.mock_llm.token_count.return_value = 100
|
||||
self.mock_llm.token_limit.return_value = 1000
|
||||
self.mock_validator = Mock(spec=XMLValidator)
|
||||
self.io_buffer = WebIOBuffer()
|
||||
self.working_memory = WorkingMemory()
|
||||
self.work_dir = Path("/tmp") # Use a temp directory for tests
|
||||
self.parser = ResponseParser(self.work_dir, self.io_buffer)
|
||||
|
||||
# Mock system metrics
|
||||
self.mock_metrics = Mock(spec=SystemMetrics)
|
||||
self.mock_metrics.get_metrics.return_value = {
|
||||
"timestamp": "2024-10-31T12:00:00Z",
|
||||
"memory_used": 1000,
|
||||
"memory_total": 2000,
|
||||
"disk_used": 5000,
|
||||
"disk_total": 10000
|
||||
}
|
||||
|
||||
# Create test agent
|
||||
self.agent = TestBaseAgent(
|
||||
system_prompt="test prompt",
|
||||
action_schema="test schema",
|
||||
working_memory=self.working_memory,
|
||||
metrics=self.mock_metrics,
|
||||
validator=self.mock_validator,
|
||||
parser=self.parser
|
||||
)
|
||||
|
||||
# Setup logging
|
||||
self.log_handler = TestLogHandler()
|
||||
logger = logging.getLogger()
|
||||
logger.addHandler(self.log_handler)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up resources."""
|
||||
logger = logging.getLogger()
|
||||
logger.removeHandler(self.log_handler)
|
||||
|
||||
if hasattr(self, 'agent'):
|
||||
del self.agent
|
||||
|
||||
def test_initialization(self):
|
||||
"""Test agent initialization and component setup."""
|
||||
self.assertEqual(self.agent._working_memory, self.working_memory)
|
||||
self.assertEqual(self.agent._metrics, self.mock_metrics)
|
||||
self.assertEqual(self.agent._validator, self.mock_validator)
|
||||
self.assertEqual(self.agent._parser, self.parser)
|
||||
self.assertEqual(self.agent._action_schema, "test schema")
|
||||
|
||||
def test_cleanup(self):
|
||||
"""Test cleanup on agent deletion."""
|
||||
# Since BaseAgent no longer has cleanup, this test is simplified
|
||||
del self.agent
|
||||
|
||||
def test_compile_context_empty_memory(self):
|
||||
"""Test context compilation with empty working memory."""
|
||||
context = self.agent._compile_context(self.mock_llm)
|
||||
|
||||
# Parse context and verify structure
|
||||
root = ET.fromstring(context)
|
||||
self.assertEqual(root.tag, "context")
|
||||
|
||||
# Check metrics
|
||||
self.assertEqual(root.get("memory_used"), "1000")
|
||||
self.assertEqual(root.get("memory_total"), "2000")
|
||||
self.assertEqual(root.get("disk_used"), "5000")
|
||||
self.assertEqual(root.get("disk_total"), "10000")
|
||||
|
||||
# Check context size
|
||||
self.assertEqual(root.get("context"), "10.0%")
|
||||
|
||||
# Check no memory entries
|
||||
self.assertEqual(len(list(root)), 0)
|
||||
|
||||
def test_compile_context_with_entries(self):
|
||||
"""Test context compilation with a single entry."""
|
||||
entry = ReasoningEntry("test-id", "test reasoning")
|
||||
self.agent._working_memory.add_entry(entry)
|
||||
|
||||
context = self.agent._compile_context(self.mock_llm)
|
||||
root = ET.fromstring(context)
|
||||
|
||||
# Check context size reflects one entry
|
||||
self.assertEqual(root.get("context"), "10.0%")
|
||||
|
||||
# Verify entry content
|
||||
reasoning_elem = root.find("reasoning")
|
||||
self.assertIsNotNone(reasoning_elem)
|
||||
self.assertEqual(reasoning_elem.get("id"), "test-id")
|
||||
self.assertEqual(reasoning_elem.text.strip(), "test reasoning")
|
||||
|
||||
def test_multiple_entries(self):
|
||||
"""Test handling multiple entries in working memory."""
|
||||
entries = [
|
||||
ReasoningEntry(f"id-{i}", f"test {i}")
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
for entry in entries:
|
||||
self.agent._working_memory.add_entry(entry)
|
||||
|
||||
context = self.agent._compile_context(self.mock_llm)
|
||||
root = ET.fromstring(context)
|
||||
|
||||
# Check context size reflects three entries
|
||||
self.assertEqual(root.get("context"), "10.0%")
|
||||
|
||||
# Verify entry order maintained
|
||||
reasoning_elems = root.findall("reasoning")
|
||||
self.assertEqual(len(reasoning_elems), 3)
|
||||
for i, elem in enumerate(reasoning_elems):
|
||||
self.assertEqual(elem.get("id"), f"id-{i}")
|
||||
self.assertEqual(elem.text.strip(), f"test {i}")
|
||||
Reference in New Issue
Block a user