162 lines
6.2 KiB
Python
162 lines
6.2 KiB
Python
import unittest
|
|
from datetime import datetime
|
|
from unittest.mock import Mock, MagicMock, patch
|
|
import xml.etree.ElementTree as ET
|
|
|
|
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.parse_error_entry import ParseErrorEntry
|
|
from sia.delete_command import DeleteCommand
|
|
from sia.entry import Entry
|
|
|
|
class TestBaseAgent(BaseAgent):
|
|
"""Concrete implementation of BaseAgent for testing."""
|
|
def run(self) -> None:
|
|
pass
|
|
|
|
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 BaseAgentTest(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test cases with mocked components."""
|
|
# Create mocks
|
|
self.mock_llm = Mock(spec=LlmEngine)
|
|
self.mock_metrics = Mock(spec=SystemMetrics)
|
|
self.mock_validator = Mock(spec=XMLValidator)
|
|
self.io_buffer = WebIOBuffer()
|
|
|
|
# Mock metrics context
|
|
metrics_elem = ET.Element("context")
|
|
metrics_elem.set("time", "2024-10-31T12:00:00Z")
|
|
metrics_elem.set("cpu", "10")
|
|
metrics_elem.set("gpu", "20")
|
|
metrics_elem.set("memory_used", "1000")
|
|
metrics_elem.set("memory_total", "2000")
|
|
metrics_elem.set("disk_used", "5000")
|
|
metrics_elem.set("disk_total", "10000")
|
|
metrics_elem.set("context", "50")
|
|
metrics_elem.set("stdin", "0")
|
|
self.mock_metrics.generate_context.return_value = metrics_elem
|
|
|
|
# Create test agent with mocked components
|
|
with patch('sia.base_agent.LlmEngine') as mock_llm_class, \
|
|
patch('sia.base_agent.SystemMetrics') as mock_metrics_class, \
|
|
patch('sia.base_agent.XMLValidator') as mock_validator_class:
|
|
|
|
mock_llm_class.return_value = self.mock_llm
|
|
mock_metrics_class.return_value = self.mock_metrics
|
|
mock_validator_class.return_value = self.mock_validator
|
|
|
|
self.agent = TestBaseAgent(
|
|
model_path="/test/model",
|
|
system_prompt="test prompt",
|
|
action_schema="test schema",
|
|
io_buffer=self.io_buffer
|
|
)
|
|
|
|
# Set timestamp for entries
|
|
self.test_timestamp = datetime(2024, 10, 31, 12, 0, 0)
|
|
|
|
def tearDown(self):
|
|
"""Clean up resources."""
|
|
if hasattr(self, 'agent'):
|
|
del self.agent
|
|
|
|
def test_initialization(self):
|
|
"""Test agent initialization and component setup."""
|
|
self.assertIsInstance(self.agent._working_memory, WorkingMemory)
|
|
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._io_buffer, self.io_buffer)
|
|
self.assertIsInstance(self.agent._parser, ResponseParser)
|
|
self.assertEqual(self.agent._system_prompt, "test prompt")
|
|
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, "state")
|
|
|
|
# Check metrics element
|
|
context_elem = root.find("context")
|
|
self.assertIsNotNone(context_elem)
|
|
self.assertEqual(context_elem.get("cpu"), "10")
|
|
self.assertEqual(context_elem.get("gpu"), "20")
|
|
|
|
# Check no memory entries
|
|
self.assertEqual(len(root.findall("*")), 1) # Only metrics element
|
|
|
|
def test_compile_context_with_entries(self):
|
|
"""Test context compilation with working memory entries."""
|
|
# Add test entry to working memory
|
|
entry = ReasoningEntry("test reasoning", "test-id", self.test_timestamp)
|
|
self.agent._working_memory.add_entry(entry)
|
|
|
|
context = self.agent._compile_context()
|
|
|
|
# Parse and verify context
|
|
root = ET.fromstring(context)
|
|
|
|
# Check metrics and entry are present
|
|
self.assertEqual(len(root.findall("*")), 2) # Metrics + reasoning
|
|
|
|
# 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, "<![CDATA[test reasoning]]>")
|
|
|
|
def test_multiple_entries(self):
|
|
"""Test handling multiple entries in working memory."""
|
|
# Add multiple entries
|
|
entries = [
|
|
ReasoningEntry(f"test {i}", f"id-{i}", self.test_timestamp)
|
|
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)
|
|
|
|
# Verify all entries present
|
|
self.assertEqual(len(root.findall("reasoning")), 3)
|
|
|
|
# Verify entry order maintained
|
|
reasoning_elems = root.findall("reasoning")
|
|
for i, elem in enumerate(reasoning_elems):
|
|
self.assertEqual(elem.get("id"), f"id-{i}")
|
|
self.assertEqual(elem.text, f"<![CDATA[test {i}]]>") |