Private instance vars, public get properties

test/background_entry_test.py

@@ -2,8 +2,7 @@ import unittest
 from datetime import datetime
 import os
 import time
-import signal
-import platform
+from typing import Optional
 
 from sia.background_entry import BackgroundEntry
 
@@ -19,7 +18,7 @@ class BackgroundEntryTest(unittest.TestCase):
         self.cleanup_files()
         
     def tearDown(self):
-        """Clean up test files and ensure no processes are left running"""
+        """Clean up test files"""
         self.cleanup_files()
         
     def cleanup_files(self):
@@ -27,53 +26,45 @@ class BackgroundEntryTest(unittest.TestCase):
         for filename in [self.test_filename, self.counter_file]:
             if os.path.exists(filename):
                 os.remove(filename)
-    
-    def wait_for_process_start(self, entry, timeout=1.0):
-        """Wait for process to start, with timeout"""
+                
+    def wait_for_condition(self, entry: BackgroundEntry, condition: callable, timeout: float = 1.0) -> bool:
+        """Wait for a condition on the entry context to be met"""
         start_time = time.time()
-        while entry.process is None and time.time() - start_time < timeout:
+        while time.time() - start_time < timeout:
             entry.update()
+            context = entry.generate_context()
+            if condition(context):
+                return True
             time.sleep(0.1)
-        self.assertIsNotNone(entry.process, "Process failed to start")
-        
-    def wait_for_process_exit(self, entry, timeout=1.0):
-        """Wait for process to exit, with timeout"""
-        start_time = time.time()
-        while entry.process is not None and time.time() - start_time < timeout:
-            entry.update()
-            time.sleep(0.1)
-        self.assertIsNone(entry.process, "Process failed to exit")
+        return False
                 
     def test_initialization(self):
         """Test entry initialization"""
-        script = "echo test"
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
+        entry = BackgroundEntry("echo test", self.test_id, self.test_timestamp)
+        context = entry.generate_context()
         
-        self.assertEqual(entry.script, script)
-        self.assertEqual(entry.stdout, "")
-        self.assertEqual(entry.stderr, "")
-        self.assertIsNone(entry.process)
-        self.assertEqual(entry.id, self.test_id)
-        self.assertEqual(entry.timestamp, self.test_timestamp)
+        # Initial context should have script but no pid or exit_code
+        self.assertEqual(context.text, "<![CDATA[echo test]]>")
+        self.assertIsNone(context.get("pid"))
+        self.assertIsNone(context.get("exit_code"))
         
     def test_short_running_process(self):
         """Test execution of a quick process"""
-        script = "echo 'test'"
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
+        entry = BackgroundEntry("echo 'test'", self.test_id, self.test_timestamp)
         
-        # Start and wait for completion
-        entry.update()
-        self.wait_for_process_start(entry)
-        self.wait_for_process_exit(entry)
+        # Wait for process to complete
+        def is_complete(ctx):
+            return ctx.get("exit_code") == "0"
+            
+        self.assertTrue(self.wait_for_condition(entry, is_complete))
         
         # Verify output
-        self.assertEqual(entry._accumulated_stdout.strip(), "test")
-        self.assertEqual(entry._accumulated_stderr, "")
-        self.assertEqual(entry._exit_code, 0)
+        context = entry.generate_context()
+        self.assertTrue("<![CDATA[test" in context.find("stdout").text)
+        self.assertEqual(context.find("stderr").text, "<![CDATA[]]>")
         
     def test_continuous_output(self):
         """Test process that generates continuous output"""
-        # Create a Python script that counts to 3 with minimal delay
         script = (
             'python3 -c "'
             'import sys\n'
@@ -85,143 +76,83 @@ class BackgroundEntryTest(unittest.TestCase):
         
         entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
         
-        # Start process and wait for completion
-        entry.update()
-        self.wait_for_process_start(entry)
-        self.wait_for_process_exit(entry, timeout=2.0)
-        
-        # One more update to ensure we get all output
-        entry.update()
-        
-        # Verify all output was captured
-        output = entry._accumulated_stdout
-        self.assertIn('count 1', output, "Missing count 1")
-        self.assertIn('count 2', output, "Missing count 2")
-        self.assertIn('count 3', output, "Missing count 3")
+        # Wait for process to complete
+        def has_all_output(ctx):
+            stdout = ctx.find("stdout").text
+            return all(f"count {i}" in stdout for i in range(1, 4))
+            
+        self.assertTrue(self.wait_for_condition(entry, has_all_output))
         
     def test_process_failure(self):
         """Test handling of process failures"""
         entry = BackgroundEntry("nonexistentcommand", self.test_id, self.test_timestamp)
         
-        # Start and wait for failure
-        entry.update()
-        self.wait_for_process_exit(entry)
-        
-        # Verify error was captured
-        self.assertTrue(len(entry._accumulated_stderr) > 0)
-        self.assertEqual(entry._accumulated_stdout, "")
-        self.assertNotEqual(entry._exit_code, 0)
-        
-    def test_generate_context_running(self):
-        """Test XML context generation for running process"""
-        script = "sleep 1"
+        # Wait for process to fail
+        def has_failed(ctx):
+            return ctx.get("exit_code") is not None and ctx.get("exit_code") != "0"
             
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
-        entry.update()  # Start process
-        self.wait_for_process_start(entry)
+        self.assertTrue(self.wait_for_condition(entry, has_failed))
         
-        element = entry.generate_context()
+        # Verify error captured
+        context = entry.generate_context()
+        self.assertTrue(len(context.find("stderr").text) > 0)
+        self.assertEqual(context.find("stdout").text, "<![CDATA[]]>")
         
-        # Verify XML structure
-        self.assertEqual(element.tag, "background")
-        self.assertEqual(element.get("id"), self.test_id)
-        self.assertIsNotNone(element.get("pid"))
-        self.assertIsNone(element.get("exit_code"))
-        
-        # Clean up
-        entry._cleanup_process()
-        
-    def test_generate_context_completed(self):
-        """Test XML context generation for completed process"""
-        script = "echo 'test'"
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
-        
-        # Start and wait for completion
-        entry.update()
-        self.wait_for_process_start(entry)
-        self.wait_for_process_exit(entry)
-        
-        element = entry.generate_context()
-        
-        # Verify XML structure
-        self.assertEqual(element.tag, "background")
-        self.assertEqual(element.get("id"), self.test_id)
-        self.assertIsNone(element.get("pid"))
-        self.assertEqual(element.get("exit_code"), "0")
-        
-        stdout_text = element.find("stdout").text
-        self.assertTrue(stdout_text.startswith("<![CDATA["))
-        self.assertIn("test", stdout_text)
-        self.assertTrue(stdout_text.endswith("]]>"))
-        
-    def test_cleanup_on_deletion(self):
-        """Test process cleanup when entry is deleted"""
-        script = "sleep 10"
-            
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
-        entry.update()
-        self.wait_for_process_start(entry)
-        
-        # Verify process is running
-        self.assertIsNotNone(entry.process)
-        pid = entry.process.pid
-        
-        # Delete the entry
-        entry._cleanup_process()
-        del entry
-        
-        # Give process time to be cleaned up
-        time.sleep(0.1)
-        
-        # Verify process was terminated
-        with self.assertRaises(ProcessLookupError):
-            os.kill(pid, 0)  # Check if process exists
-
     def test_cleanup_running_process(self):
         """Test cleanup of running process"""
-        script = "sleep 10"
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
+        entry = BackgroundEntry("sleep 10", self.test_id, self.test_timestamp)
         
-        # Start the process
-        entry.update()
-        self.assertIsNotNone(entry.process)
-        process_pid = entry.process.pid
+        # Wait for process to start
+        def has_started(ctx):
+            return ctx.get("pid") is not None
         
-        # Cleanup should terminate process
+        self.assertTrue(self.wait_for_condition(entry, has_started))
+        
+        # Get PID before cleanup
+        context = entry.generate_context()
+        pid = int(context.get("pid"))
+        
+        # Clean up and verify process terminated
         entry.cleanup()
         
-        # Verify process was terminated
-        self.assertIsNone(entry.process)
+        # Process should no longer exist
         time.sleep(0.1)  # Give process time to terminate
         with self.assertRaises(ProcessLookupError):
-            import os
-            os.kill(process_pid, 0)
+            os.kill(pid, 0)
             
     def test_cleanup_completed_process(self):
         """Test cleanup of already completed process"""
-        script = "echo 'test'"
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
+        entry = BackgroundEntry("echo test", self.test_id, self.test_timestamp)
         
-        # Run and wait for completion
-        entry.update()
-        self.wait_for_process_exit(entry)
+        # Wait for completion
+        def is_complete(ctx):
+            return ctx.get("exit_code") == "0"
+            
+        self.assertTrue(self.wait_for_condition(entry, is_complete))
         
-        # Cleanup should be safe for completed process
+        # Cleanup should be safe
         entry.cleanup()
-        self.assertIsNone(entry.process)
         
     def test_multiple_cleanup_calls(self):
         """Test that multiple cleanup calls are safe"""
-        script = "sleep 10"
-        entry = BackgroundEntry(script, self.test_id, self.test_timestamp)
+        entry = BackgroundEntry("sleep 10", self.test_id, self.test_timestamp)
         
-        # Start process
-        entry.update()
-        self.assertIsNotNone(entry.process)
+        # Wait for process to start
+        def has_started(ctx):
+            return ctx.get("pid") is not None
+        
+        self.assertTrue(self.wait_for_condition(entry, has_started))
+        
+        # Get PID before cleanup
+        context = entry.generate_context()
+        pid = int(context.get("pid"))
         
         # Multiple cleanups should be safe
         entry.cleanup()
         entry.cleanup()
         entry.cleanup()
         
-        self.assertIsNone(entry.process)
+        # Process should no longer exist
+        time.sleep(0.1)
+        with self.assertRaises(ProcessLookupError):
+            os.kill(pid, 0)

test/delete_command_test.py

@@ -1,18 +1,59 @@
-import unittest
 from datetime import datetime
+import os
 import time
+import unittest
 
 from sia.working_memory import WorkingMemory
 from sia.delete_command import DeleteCommand
 from sia.reasoning_entry import ReasoningEntry
 from sia.background_entry import BackgroundEntry
 
+
+
 class DeleteCommandTest(unittest.TestCase):
     def setUp(self):
         """Set up test cases with a fresh working memory"""
         self.memory = WorkingMemory()
         self.test_id = "test-id-1234"
         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_delete_running_background_entry(self):
+        """Test deleting a background entry that is still running"""
+        # Create and start background process
+        entry = BackgroundEntry("sleep 10", "test-id", self.test_timestamp)
+        self.memory.add_entry(entry)
+        
+        # Wait for process to start
+        self.assertTrue(self.wait_for_process_start(entry))
+        
+        # Get PID before deletion
+        context = entry.generate_context()
+        pid = int(context.get("pid"))
+        
+        # Delete entry and verify process terminated
+        command = DeleteCommand(entry.id)
+        result = command.execute(self.memory)
+        
+        self.assertTrue(result.success)
+        self.assertFalse(result.should_stop)
+        self.assertEqual(result.message, "")
+        self.assertEqual(self.memory.get_entries_count(), 0)
+        
+        # Verify process was terminated
+        time.sleep(0.1)  # Give process time to terminate
+        with self.assertRaises(ProcessLookupError):
+            os.kill(pid, 0)
         
     def test_delete_existing_entry(self):
         """Test deleting an existing entry"""
@@ -29,31 +70,6 @@ class DeleteCommandTest(unittest.TestCase):
         self.assertFalse(result.should_stop)
         self.assertEqual(result.message, "")
         self.assertEqual(self.memory.get_entries_count(), 0)
-        
-    def test_delete_running_background_entry(self):
-        """Test deleting a background entry that is still running"""
-        # Create a background entry with a long-running command
-        entry = BackgroundEntry("sleep 10", self.test_id, self.test_timestamp)
-        self.memory.add_entry(entry)
-        
-        # Start the process
-        entry.update()
-        self.assertIsNotNone(entry.process)
-        process_pid = entry.process.pid
-        
-        # Execute delete command
-        command = DeleteCommand(self.test_id)
-        result = command.execute(self.memory)
-        
-        # Verify process was terminated
-        time.sleep(0.1)  # Give process time to terminate
-        self.assertTrue(result.success)
-        self.assertIsNone(entry.process)
-        
-        # Verify process no longer exists
-        with self.assertRaises(ProcessLookupError):
-            import os
-            os.kill(process_pid, 0)
             
     def test_delete_nonexistent_entry(self):
         """Test attempting to delete a nonexistent entry"""

test/llm_engine_test.py

@@ -0,0 +1,25 @@
+import unittest
+
+from itertools import tee
+
+from . import test_data
+
+from sia.llm_engine import LlmEngine
+
+class LlmEngineTest(unittest.TestCase):
+    def setUp(self):
+        self.model_path = "/root/model"
+
+    def test_initialization(self):
+        llm_engine = LlmEngine(self.model_path)
+        self.assertIsInstance(llm_engine, LlmEngine)
+
+    def test_infer(self):
+        main_context = "This is a test"
+        llm_engine = LlmEngine(self.model_path)
+        tokens = llm_engine.infer(test_data.echo_system_prompt, main_context)
+        print_tokens, result_tokens = tee(tokens)
+        for token in print_tokens:
+            print(token, end="", flush=True)
+        result = ''.join(result_tokens)
+        self.assertEqual(result, f"{main_context}<test_tag>{main_context}</test_tag>")

test/parse_error_entry_test.py

@@ -15,8 +15,6 @@ class ParseErrorEntryTest(unittest.TestCase):
         error = "parsing failed"
         entry = ParseErrorEntry(content, error, self.test_id, self.test_timestamp)
         
-        self.assertEqual(entry.content, content)
-        self.assertEqual(entry.error, error)
         self.assertEqual(entry.id, self.test_id)
         self.assertEqual(entry.timestamp, self.test_timestamp)
         
@@ -27,15 +25,13 @@ class ParseErrorEntryTest(unittest.TestCase):
         entry = ParseErrorEntry(content, error, self.test_id, self.test_timestamp)
         
         # Store initial state
-        initial_content = entry.content
-        initial_error = entry.error
+        initial_context = str(entry.generate_context())
         
         # Perform update
         entry.update()
         
         # Verify state hasn't changed
-        self.assertEqual(entry.content, initial_content)
-        self.assertEqual(entry.error, initial_error)
+        self.assertEqual(str(entry.generate_context()), initial_context)
         
     def test_generate_context(self):
         """Test XML context generation"""
@@ -75,14 +71,4 @@ class ParseErrorEntryTest(unittest.TestCase):
         
         element = entry.generate_context()
         self.assertEqual(element.find("content").text, "<![CDATA[]]>")
-        self.assertEqual(element.find("error").text, "<![CDATA[]]>")
-        
-    def test_large_content(self):
-        """Test handling large content and error messages"""
-        content = "x" * 10000
-        error = "y" * 10000
-        entry = ParseErrorEntry(content, error, self.test_id, self.test_timestamp)
-        
-        element = entry.generate_context()
-        self.assertEqual(element.find("content").text, f"<![CDATA[{content}]]>")
-        self.assertEqual(element.find("error").text, f"<![CDATA[{error}]]>")
+        self.assertEqual(element.find("error").text, "<![CDATA[]]>")

test/read_entry_test.py

@@ -18,7 +18,6 @@ class ReadEntryTest(unittest.TestCase):
         self.assertEqual(entry.content, "")
         self.assertEqual(entry.id, self.test_id)
         self.assertEqual(entry.timestamp, self.test_timestamp)
-        self.assertFalse(entry._read)
         
     def test_single_read(self):
         """Test reading content once"""
@@ -28,10 +27,8 @@ class ReadEntryTest(unittest.TestCase):
         entry = ReadEntry(self.io_buffer, self.test_id, self.test_timestamp)
         entry.update()
         
-        # Verify content was read
         self.assertEqual(entry.content, test_input)
-        self.assertTrue(entry._read)
-        self.assertEqual(self.io_buffer.buffer_length(), 0)  # Buffer should be cleared
+        self.assertEqual(self.io_buffer.buffer_length(), 0)
         
     def test_multiple_updates(self):
         """Test that content is only read once even with multiple updates"""
@@ -42,14 +39,11 @@ class ReadEntryTest(unittest.TestCase):
         entry.update()
         initial_content = entry.content
         
-        # Add more input and update again
         self.io_buffer.append_stdin("additional input")
         entry.update()
         entry.update()
         
-        # Verify content hasn't changed after first read
         self.assertEqual(entry.content, initial_content)
-        self.assertTrue(entry._read)
         
     def test_empty_input(self):
         """Test reading when no input is available"""
@@ -57,7 +51,6 @@ class ReadEntryTest(unittest.TestCase):
         entry.update()
         
         self.assertEqual(entry.content, "")
-        self.assertTrue(entry._read)
         
     def test_special_characters(self):
         """Test reading content with special characters"""

test/reasoning_entry_test.py

@@ -24,13 +24,13 @@ class ReasoningEntryTest(unittest.TestCase):
         entry = ReasoningEntry(content, self.test_id, self.test_timestamp)
         
         # Store initial state
-        initial_content = entry.content
+        initial_content = entry.generate_context().text
         
         # Perform update
         entry.update()
         
         # Verify state hasn't changed
-        self.assertEqual(entry.content, initial_content)
+        self.assertEqual(entry.generate_context().text, initial_content)
         
     def test_generate_context(self):
         """Test XML context generation"""

test/repeat_entry_test.py

@@ -145,20 +145,4 @@ class RepeatEntryTest(unittest.TestCase):
         outputs = [first_output, second_output, third_output]
         unique_outputs = set(outputs)
         self.assertEqual(len(unique_outputs), 3)
-        self.assertEqual(outputs, ['Output 1', 'Output 2', 'Output 3'])
-        
-    def test_error_recovery(self):
-        """Test that entry recovers from errors in subsequent updates"""
-        # Start with a failing command
-        entry = RepeatEntry("nonexistentcommand", self.test_id, self.test_timestamp)
-        entry.update()
-        self.assertNotEqual(entry.exit_code, 0)
-        
-        # Change to a working command
-        entry.script = "echo 'test'"
-        entry.update()
-        
-        # Verify recovery
-        self.assertEqual(entry.exit_code, 0)
-        self.assertEqual(entry.stdout.strip(), "test")
-        self.assertEqual(entry.stderr, "")
+        self.assertEqual(outputs, ['Output 1', 'Output 2', 'Output 3'])

test/response_parser_test.py

@@ -52,8 +52,8 @@ class ResponseParserTest(unittest.TestCase):
         result = self.parser.parse(xml)
         
         self.assertIsInstance(result, BackgroundEntry)
-        self.assertEqual(result.script, "echo test")
-        
+        self.assertEqual(result.script, "echo test")        
+
     def test_background_entry_empty_script(self):
         """Test parsing background entry with empty script returns error entry"""
         xml = '<background/>'
@@ -120,7 +120,6 @@ class ResponseParserTest(unittest.TestCase):
         result = self.parser.parse(xml)
         
         self.assertIsInstance(result, ReadEntry)
-        self.assertEqual(result.io_buffer, self.io_buffer)
         
     def test_write_stdout_entry(self):
         """Test parsing write stdout entry"""

test/stop_command_test.py

@@ -1,10 +1,10 @@
-import unittest
 from datetime import datetime
+import os
 import time
+import unittest
 
 from sia.working_memory import WorkingMemory
 from sia.stop_command import StopCommand
-from sia.reasoning_entry import ReasoningEntry
 from sia.background_entry import BackgroundEntry
 
 class StopCommandTest(unittest.TestCase):
@@ -12,34 +12,38 @@ class StopCommandTest(unittest.TestCase):
         """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_stop_command_cleanup(self):
         """Test stop command cleans up and removes all entries"""
-        # Add a mix of entry types including a background process
-        entries = [
-            ReasoningEntry("test content", "id-1", self.test_timestamp),
-            BackgroundEntry("sleep 10", "id-2", self.test_timestamp)
-        ]
+        # Add background process
+        entry = BackgroundEntry("sleep 10", "test-id", self.test_timestamp)
+        self.memory.add_entry(entry)
         
-        for entry in entries:
-            self.memory.add_entry(entry)
-            
-        # Start the background process
-        background_entry = self.memory.get_entry("id-2")
-        background_entry.update()
-        self.assertIsNotNone(background_entry.process)
-        process_pid = background_entry.process.pid
+        # 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"))
         
         # Execute stop command
         command = StopCommand()
         result = command.execute(self.memory)
         
-        # Verify all entries removed
+        # Verify entries removed and process terminated
         self.assertTrue(result.should_stop)
         self.assertEqual(self.memory.get_entries_count(), 0)
         
-        # Verify background process was terminated
+        # Verify process was terminated
         time.sleep(0.1)  # Give process time to terminate
         with self.assertRaises(ProcessLookupError):
-            import os
-            os.kill(process_pid, 0)
+            os.kill(pid, 0)

test/working_memory_test.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+import os
 import time
 import unittest
 import xml.etree.ElementTree as ET
@@ -12,6 +13,7 @@ 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):
@@ -31,6 +33,17 @@ class WorkingMemoryTest(unittest.TestCase):
         """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"""
@@ -147,46 +160,40 @@ class WorkingMemoryTest(unittest.TestCase):
 
     def test_remove_entry_cleanup(self):
         """Test that removing an entry triggers cleanup"""
-        # Create a background process entry
-        entry = BackgroundEntry("sleep 10", "test-id", datetime(2024, 10, 31, 12, 0, 0))
+        # Create and start background process
+        entry = BackgroundEntry("sleep 10", "test-id", self.test_timestamp)
         self.memory.add_entry(entry)
         
-        # Start the process
-        entry.update()
-        self.assertIsNotNone(entry.process)
-        process_pid = entry.process.pid
+        # 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 should trigger cleanup
-        self.memory.remove_entry("test-id")
+        # Remove entry and verify process terminated
+        self.memory.remove_entry(entry.id)
         
-        # Verify process was terminated
         time.sleep(0.1)  # Give process time to terminate
         with self.assertRaises(ProcessLookupError):
-            import os
-            os.kill(process_pid, 0)
-            
+            os.kill(pid, 0)
+
     def test_cleanup_on_memory_clear(self):
         """Test that clearing memory properly cleans up all entries"""
         # Add multiple background processes
-        processes = []
+        pids = []
         for i in range(3):
-            entry = BackgroundEntry(
-                "sleep 10",
-                f"id-{i}",
-                datetime(2024, 10, 31, 12, 0, 0)
-            )
+            entry = BackgroundEntry("sleep 10", f"id-{i}", self.test_timestamp)
             self.memory.add_entry(entry)
-            entry.update()
-            processes.append(entry.process.pid)
+            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 processes:
+        for pid in pids:
             with self.assertRaises(ProcessLookupError):
-                import os
                 os.kill(pid, 0)
         
         self.assertEqual(self.memory.get_entries_count(), 0)
@@ -194,14 +201,13 @@ class WorkingMemoryTest(unittest.TestCase):
     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",
-            datetime(2024, 10, 31, 12, 0, 0)
-        )
+        entry = BackgroundEntry("sleep 10", "test-id", self.test_timestamp)
         self.memory.add_entry(entry)
-        entry.update()
-        process_pid = entry.process.pid
+        
+        # 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
@@ -209,8 +215,7 @@ class WorkingMemoryTest(unittest.TestCase):
         # Verify process was terminated
         time.sleep(0.1)  # Give process time to terminate
         with self.assertRaises(ProcessLookupError):
-            import os
-            os.kill(process_pid, 0)
+            os.kill(pid, 0)
 
     def test_get_entries(self):
         """Test getting all entries"""