Start work on docker module

test/docker_module_test.py

@@ -0,0 +1,186 @@
+import unittest
+import docker
+from docker.models.containers import Container
+from sia.docker_module import DockerModule
+
+class DockerModuleTest(unittest.TestCase):
+    def test_initialization(self):
+        """Test DockerModule initialization."""
+        docker_module = DockerModule()
+        self.assertIsInstance(docker_module, DockerModule)
+
+    def test_start_container_short_lived(self):
+        """Test starting a short-lived container."""
+        docker_module = DockerModule()
+        
+        # Test parameters
+        image = "sia:latest"
+        timeout = 1000
+        command = "echo"
+        arguments = ["Hello World"]
+        
+        output = docker_module.start_container(
+            image=image,
+            timeout=timeout,
+            command=command,
+            arguments=arguments
+        )
+        
+        # Verify output was returned
+        self.assertEqual(output, "Hello World\n")
+
+    def test_start_container_long_running(self):
+        """Test starting a long-running container."""
+        docker_module = DockerModule()
+        
+        name = "test-container"
+        image = "sia:latest"
+        
+        container_name = docker_module.start_container(
+            image=image,
+            name=name
+        )
+        
+        # Verify container was started and stored
+        self.assertEqual(container_name, name)
+        self.assertIn(name, docker_module.containers)
+        self.assertEqual(docker_module.containers[name], self.mock_container)
+
+    def test_start_container_with_volumes(self):
+        """Test starting a container with volume mappings."""
+        docker_module = DockerModule()
+        
+        volumes = {"./pdf": "/pdf"}
+        output = docker_module.start_container(
+            image="sia:latest",
+            timeout=1000,
+            volumes=volumes
+        )
+        
+        # Verify volume bindings were configured correctly
+        expected_binds = {
+            "./pdf": {"bind": "/pdf", "mode": "rw"}
+        }
+
+    def test_invalid_container_start(self):
+        """Test error handling for invalid container start."""
+        docker_module = DockerModule()
+        
+        # Test missing required parameters
+        with self.assertRaises(ValueError):
+            docker_module.start_container(image="test:latest")
+
+    def test_stop_container(self):
+        """Test stopping a container."""
+        docker_module = DockerModule()
+        
+        # Start a container first
+        name = "test-container"
+        docker_module.start_container(
+            image="test:latest",
+            name=name
+        )
+        
+        # Stop the container
+        docker_module.stop_container(name)
+        
+        # Verify container was stopped and removed
+        self.mock_container.stop.assert_called_once()
+        self.mock_container.remove.assert_called_once()
+        self.assertNotIn(name, docker_module.containers)
+
+    def test_stop_nonexistent_container(self):
+        """Test stopping a container that doesn't exist."""
+        docker_module = DockerModule()
+        
+        with self.assertRaises(KeyError):
+            docker_module.stop_container("nonexistent")
+
+    def test_container_io_operations(self):
+        """Test container I/O operations."""
+        docker_module = DockerModule()
+        
+        # Start a container
+        name = "test-container"
+        docker_module.start_container(
+            image="test:latest",
+            name=name
+        )
+        
+        # Test stdin write
+        test_input = "test input"
+        mock_socket = Mock()
+        mock_socket._sock = Mock()
+        self.mock_container.attach_socket.return_value = mock_socket
+        
+        docker_module.write_container_stdin(name, test_input)
+        mock_socket._sock.send.assert_called_once_with(test_input.encode('utf-8'))
+        
+        # Test stdout read
+        stdout = docker_module.read_container_stdout(name)
+        self.mock_container.logs.assert_called_with(
+            stdout=True, 
+            stderr=False, 
+            tail='all'
+        )
+        self.assertEqual(stdout, "Test output")
+        
+        # Test stderr read
+        stderr = docker_module.read_container_stderr(name)
+        self.mock_container.logs.assert_called_with(
+            stdout=False, 
+            stderr=True, 
+            tail='all'
+        )
+        self.assertEqual(stderr, "Test output")
+
+    def test_wait_container(self):
+        """Test waiting for container completion."""
+        docker_module = DockerModule()
+        
+        # Start a container
+        name = "test-container"
+        docker_module.start_container(
+            image="test:latest",
+            name=name
+        )
+        
+        # Test wait operation
+        exit_code, output = docker_module.wait_container(name, timeout=1000)
+        
+        # Verify wait was called with correct timeout
+        self.mock_container.wait.assert_called_once_with(timeout=1.0)
+        self.assertEqual(exit_code, 0)
+        self.assertEqual(output, "Test output")
+
+    def test_get_container_status(self):
+        """Test getting container status information."""
+        docker_module = DockerModule()
+        
+        # Start a container
+        name = "test-container"
+        docker_module.start_container(
+            image="test:latest",
+            name=name
+        )
+        
+        # Configure mock logs sizes
+        self.mock_container.logs.side_effect = [b"stdout", b"stderr"]
+        
+        # Get status
+        status = docker_module.get_container_status(name)
+        
+        # Verify status information
+        self.assertEqual(status['name'], name)
+        self.assertEqual(status['status'], "running")
+        self.assertEqual(status['started_at'], "2024-10-23T10:00:00Z")
+        self.assertEqual(status['exit_code'], 0)
+        self.assertEqual(status['error'], "")
+        self.assertEqual(status['stdout_size'], 6)  # len(b"stdout")
+        self.assertEqual(status['stderr_size'], 6)  # len(b"stderr")
+        
+        # Verify container was reloaded
+        self.mock_container.reload.assert_called_once()
+
+if __name__ == '__main__':
+    unittest.main()

test/llm_engine_test.py

@@ -1,14 +1,15 @@
 import unittest
 
+from itertools import tee
+
 from . import test_data
 from . import test_util
 
-from sia.llm_engine import LlmEngine, InferenceResult
+from sia.llm_engine import LlmEngine
 
 class LlmEngineTest(unittest.TestCase):
     def setUp(self):
         self.model_path = "/root/model"
-        self.llm_engine = LlmEngine(self.model_path)
 
     def test_initialization(self):
         llm_engine = LlmEngine(self.model_path)
@@ -17,9 +18,9 @@ class LlmEngineTest(unittest.TestCase):
     def test_infer(self):
         main_context = "This is a test"
         llm_engine = LlmEngine(self.model_path)
-        result = llm_engine.infer(test_data.echo_system_prompt, main_context, test_data.echo_action_schema)
-        self.assertIsInstance(result, InferenceResult)
-        self.assertIsInstance(result.reasoning, str)
-        self.assertIsInstance(result.actions, str)
-        self.assertEqual(result.reasoning, main_context)
-        self.assertEqual(result.actions, f"<test_tag>{main_context}</test_tag>")
+        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>")