from datetime import datetime import os import time import unittest import xml.etree.ElementTree as ET from sia.background_entry import BackgroundEntry from sia.entry import Entry from sia.parse_error_entry import ParseErrorEntry from sia.read_entry import ReadEntry from sia.reasoning_entry import ReasoningEntry from sia.web_io_buffer import WebIOBuffer from sia.working_memory import WorkingMemory from sia.write_entry import WriteEntry class MockEntry(Entry): """Mock entry class for testing""" def __init__(self, id: str, timestamp: datetime): super().__init__(id, timestamp) self.updated = False def update(self) -> None: self.updated = True def generate_context(self) -> ET.Element: elem = ET.Element("mock", {"id": self.id}) elem.text = "" return elem class WorkingMemoryTest(unittest.TestCase): def setUp(self): """Set up test cases with a fresh working memory""" self.memory = WorkingMemory() self.test_timestamp = datetime(2024, 10, 31, 12, 0, 0) def wait_for_process_start(self, entry: BackgroundEntry, timeout=1.0): """Wait for process to start, checking context pid""" start_time = time.time() while time.time() - start_time < timeout: entry.update() context = entry.generate_context() if context.get("pid") is not None: return True time.sleep(0.1) return False def test_initialization(self): """Test initial state of working memory""" self.assertEqual(self.memory.get_entries_count(), 0) self.assertEqual(self.memory.generate_context(), []) def test_add_entry(self): """Test adding entries""" entry = MockEntry("test-id-1", self.test_timestamp) self.memory.add_entry(entry) self.assertEqual(self.memory.get_entries_count(), 1) self.assertEqual(self.memory.get_entry("test-id-1"), entry) def test_add_invalid_entry(self): """Test adding invalid entry type""" with self.assertRaises(TypeError): self.memory.add_entry("not an entry") def test_remove_entry(self): """Test removing entries""" entry1 = MockEntry("test-id-1", self.test_timestamp) entry2 = MockEntry("test-id-2", self.test_timestamp) self.memory.add_entry(entry1) self.memory.add_entry(entry2) self.memory.remove_entry("test-id-1") self.assertEqual(self.memory.get_entries_count(), 1) self.assertIsNone(self.memory.get_entry("test-id-1")) self.assertEqual(self.memory.get_entry("test-id-2"), entry2) def test_remove_nonexistent_entry(self): """Test removing entry that doesn't exist""" entry = MockEntry("test-id-1", self.test_timestamp) self.memory.add_entry(entry) self.memory.remove_entry("nonexistent-id") self.assertEqual(self.memory.get_entries_count(), 1) self.assertEqual(self.memory.get_entry("test-id-1"), entry) def test_update_entries(self): """Test updating all entries""" entries = [ MockEntry(f"test-id-{i}", self.test_timestamp) for i in range(3) ] for entry in entries: self.memory.add_entry(entry) self.memory.update() for entry in entries: self.assertTrue(entry.updated) def test_generate_context(self): """Test generating XML context""" entry1 = MockEntry("test-id-1", self.test_timestamp) entry2 = MockEntry("test-id-2", self.test_timestamp) self.memory.add_entry(entry1) self.memory.add_entry(entry2) context = self.memory.generate_context() self.assertEqual(len(context), 2) self.assertEqual(context[0].tag, "mock") self.assertEqual(context[0].get("id"), "test-id-1") self.assertEqual(context[1].tag, "mock") self.assertEqual(context[1].get("id"), "test-id-2") def test_get_entries_by_type(self): """Test retrieving entries by type""" # Create IO buffer for IO entries io_buffer = WebIOBuffer() # Add different types of entries entries = [ ReasoningEntry("test reasoning", "reasoning-1", self.test_timestamp), ParseErrorEntry("bad content", "error msg", "error-1", self.test_timestamp), ReadEntry(io_buffer, "read-1", self.test_timestamp), WriteEntry("test output", io_buffer, "write-1", self.test_timestamp) ] for entry in entries: self.memory.add_entry(entry) # Test filtering by each type reasoning_entries = self.memory.get_entries_by_type(ReasoningEntry) self.assertEqual(len(reasoning_entries), 1) self.assertIsInstance(reasoning_entries[0], ReasoningEntry) error_entries = self.memory.get_entries_by_type(ParseErrorEntry) self.assertEqual(len(error_entries), 1) self.assertIsInstance(error_entries[0], ParseErrorEntry) read_entries = self.memory.get_entries_by_type(ReadEntry) self.assertEqual(len(read_entries), 1) self.assertIsInstance(read_entries[0], ReadEntry) write_entries = self.memory.get_entries_by_type(WriteEntry) self.assertEqual(len(write_entries), 1) self.assertIsInstance(write_entries[0], WriteEntry) def test_empty_context_generation(self): """Test context generation with no entries""" context = self.memory.generate_context() self.assertEqual(context, []) def test_get_nonexistent_entry(self): """Test retrieving entry that doesn't exist""" self.assertIsNone(self.memory.get_entry("nonexistent-id")) def test_remove_entry_cleanup(self): """Test that removing an entry triggers cleanup""" # Create and start background process entry = BackgroundEntry("sleep 10", "test-id", self.test_timestamp) self.memory.add_entry(entry) # Wait for process to start and get PID self.assertTrue(self.wait_for_process_start(entry)) context = entry.generate_context() pid = int(context.get("pid")) # Remove entry and verify process terminated self.memory.remove_entry(entry.id) time.sleep(0.1) # Give process time to terminate with self.assertRaises(ProcessLookupError): os.kill(pid, 0) def test_cleanup_on_memory_clear(self): """Test that clearing memory properly cleans up all entries""" # Add multiple background processes pids = [] for i in range(3): entry = BackgroundEntry("sleep 10", f"id-{i}", self.test_timestamp) self.memory.add_entry(entry) self.assertTrue(self.wait_for_process_start(entry)) context = entry.generate_context() pids.append(int(context.get("pid"))) # Clear memory self.memory.clear() # Verify all processes were terminated time.sleep(0.1) # Give processes time to terminate for pid in pids: with self.assertRaises(ProcessLookupError): os.kill(pid, 0) self.assertEqual(self.memory.get_entries_count(), 0) def test_cleanup_on_memory_deletion(self): """Test that deleting memory properly cleans up all entries""" # Add a background process entry = BackgroundEntry("sleep 10", "test-id", self.test_timestamp) self.memory.add_entry(entry) # Wait for process to start and get PID self.assertTrue(self.wait_for_process_start(entry)) context = entry.generate_context() pid = int(context.get("pid")) # Delete memory del self.memory # Verify process was terminated time.sleep(0.1) # Give process time to terminate with self.assertRaises(ProcessLookupError): os.kill(pid, 0) def test_get_entries(self): """Test getting all entries""" # Add multiple entries entries = [ ReasoningEntry(f"content {i}", f"id-{i}", self.test_timestamp) for i in range(3) ] for entry in entries: self.memory.add_entry(entry) # Get all entries retrieved_entries = self.memory.get_entries() # Verify count and contents self.assertEqual(len(retrieved_entries), len(entries)) for entry in entries: self.assertIn(entry, retrieved_entries) def test_get_entries_returns_copy(self): """Test that get_entries returns a copy of the list""" # Add an entry entry = ReasoningEntry("test content", "test-id", self.test_timestamp) self.memory.add_entry(entry) # Get entries and modify the returned list entries = self.memory.get_entries() entries.clear() # Verify original memory is unchanged self.assertEqual(self.memory.get_entries_count(), 1) self.assertIsNotNone(self.memory.get_entry("test-id"))