Docker module basic functionality

sia/docker_module.py

@@ -2,6 +2,7 @@ from docker.models.containers import Container
 from typing import Dict, List, Optional, Tuple
 import docker
 import logging
+import requests
 import subprocess
 import time
 
@@ -21,10 +22,54 @@ class DockerModule:
     def __init__(self):
         DockerModule._start_docker_engine()
         self.client = docker.DockerClient(base_url='unix://var/run/docker.sock')
-        self.containers: Dict[str, Container] = {}
+        self.containers: Dict[str, (Container, docker.APIClient.AttachSocket)] = {}
-        self.container_sockets: Dict[str, docker.APIClient.AttachSocket] = {}
         self.logger = logging.getLogger(__name__)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exec_type, exec_value, traceback):
+        del self
+    
+    def __del__(self):
+        """Cleanup all resources before shutdown."""
+        try:
+            # Stop and remove all containers
+            for (container, socket) in self.containers.values():
+                self._wait_and_stop(container, socket, 0)
+            # Close the Docker client
+            if hasattr(self, 'client'):
+                self.client.close()
+        except Exception as e:
+            self.logger.error(f"Error during cleanup: {str(e)}")
+
+    def _wait_and_stop(self, container: Container, socket, timeout: int) -> Tuple[int, str]:
+        """Helper method for waiting on long- and short-lived containers."""
+        status_code = None
+        if timeout > 0:
+            try:
+                result = container.wait(timeout=timeout/1000)
+                status_code = result.get('StatusCode')
+            except requests.exceptions.ReadTimeout as e:
+                container.stop()
+                status_code = -1
+            except requests.exceptions.ConnectionError as e:
+                container.stop()
+                status_code = -1
+        else:
+            container.stop()
+            status_code = -1
+        try:
+            socket.close()
+        except Exception as e:
+            self.logger.error(f"Error closing socket: {str(e)}")
+        output = container.logs(stdout=True, stderr=True).decode('utf-8')
+        try:
+            container.remove(force=True)
+        except Exception as e:
+            self.logger.error(f"Error removing container: {str(e)}")
+        return status_code, output
+
     def start_container(
         self,
         image: str,
@@ -59,100 +104,46 @@ class DockerModule:
         # Validate inputs
         if name is None and timeout < 0:
             raise ValueError("Either name or timeout must be provided")
-
-        # Prepare command with arguments
         cmd = []
         if command:
             cmd.append(command)
         if arguments:
             cmd.extend(arguments)
-
-        # Prepare volume bindings
         volume_binds = {}
         if volumes:
             for host_path, container_path in volumes.items():
                 volume_binds[host_path] = {'bind': container_path, 'mode': 'rw'}
-
-        # Prepare port bindings
         port_bindings = {}
         if ports:
             for host_port, container_port in ports.items():
                 port_bindings[container_port] = host_port
-
+        self.client.images.pull(image)
-        try:
+        container = self.client.containers.create(
-            self.client.images.pull('busybox:latest')
+            image,
-            # Create container with tty and open stdin
+            command=cmd if cmd else None,
-            container = self.client.containers.create(
+            name=name,
-                image,
+            detach=True,
-                command=cmd if cmd else None,
+            volumes=volume_binds,
-                name=name,
+            ports=port_bindings,
-                detach=True,
+            environment=environment,
-                volumes=volume_binds,
+            tty=False,
-                ports=port_bindings,
+            stdin_open=True
-                environment=environment,
+        )
-                tty=True,
+        socket = container.attach_socket(params={
-                stdin_open=True
+            'stdin': 1,
-            )
+            'stdout': 1,
-
+            'stderr': 1,
-            # Attach socket before starting container
+            'stream': 1,
-            socket = container.attach_socket(params={
+            'demux': True
-                'stdin': 1,
+        })
-                'stdout': 1,
+        container.start()
-                'stderr': 1,
+        if timeout >= 0:
-                'stream': 1
+            # For short-lived containers
-            })
+            return self._wait_and_stop(container, socket, timeout)[1]
-
+        else:
-            # Start the container
+            # For long-running containers
-            container.start()
+            self.containers[name] = (container, socket)
-
+            return name
-            if timeout >= 0:
-                # For short-lived containers, wait for completion
-                try:
-                    container.wait(timeout=timeout/1000)  # Convert ms to seconds
-                    output = container.logs().decode('utf-8')
-                    socket.close()
-                    container.remove(force=True)
-                    return output
-                except docker.errors.NotFound:
-                    self.logger.warning(f"Container was removed before timeout")
-                    return ""
-            else:
-                # For long-running containers, store references
-                self.containers[name] = container
-                self.container_sockets[name] = socket
-                return name
-
-        except Exception as e:
-            self.logger.error(f"Error starting container: {str(e)}")
-            raise
-
-    def stop_container(self, name: str) -> None:
-        """Stop and remove a Docker container by name.
-
-        Args:
-            name: Name of the container to stop
-
-        Raises:
-            KeyError: If container name not found
-            docker.errors.APIError: If container stop/removal fails
-        """
-        if name not in self.containers:
-            raise KeyError(f"Container {name} not found")
-
-        try:
-            # Close socket if it exists
-            if name in self.container_sockets:
-                self.container_sockets[name].close()
-                del self.container_sockets[name]
-
-            # Stop and remove container
-            container = self.containers[name]
-            container.stop()
-            container.remove()
-            del self.containers[name]
-        except Exception as e:
-            self.logger.error(f"Error stopping container {name}: {str(e)}")
-            raise
 
     def write_container_stdin(self, name: str, data: str) -> None:
         """Write data to a container's standard input.
@@ -165,11 +156,11 @@ class DockerModule:
             KeyError: If container name not found
             docker.errors.APIError: If write operation fails
         """
-        if name not in self.containers or name not in self.container_sockets:
+        if name not in self.containers:
-            raise KeyError(f"Container {name} not found or socket not initialized")
+            raise KeyError(f"Container {name} not found")
 
         try:
-            socket = self.container_sockets[name]
+            _, socket = self.containers[name]
             socket._sock.send(data.encode('utf-8'))
         except Exception as e:
             self.logger.error(f"Error writing to container {name} stdin: {str(e)}")
@@ -193,7 +184,7 @@ class DockerModule:
             raise KeyError(f"Container {name} not found")
 
         try:
-            container = self.containers[name]
+            container, socket = self.containers[name]
             logs = container.logs(stdout=True, stderr=False, tail=n if n > 0 else 'all')
             return logs.decode('utf-8')
         except Exception as e:
@@ -218,7 +209,7 @@ class DockerModule:
             raise KeyError(f"Container {name} not found")
 
         try:
-            container = self.containers[name]
+            container, socket = self.containers[name]
             logs = container.logs(stdout=False, stderr=True, tail=n if n > 0 else 'all')
             return logs.decode('utf-8')
         except Exception as e:
@@ -243,14 +234,8 @@ class DockerModule:
         if name not in self.containers:
             raise KeyError(f"Container {name} not found")
 
-        try:
+        container, socket = self.containers.pop(name)
-            container = self.containers[name]
+        return self._wait_and_stop(container, socket, timeout)
-            result = container.wait(timeout=timeout/1000)  # Convert ms to seconds
-            logs = container.logs().decode('utf-8')
-            return (result['StatusCode'], logs)
-        except Exception as e:
-            self.logger.error(f"Error waiting for container {name}: {str(e)}")
-            raise
 
     def get_container_status(self, name: str) -> Dict[str, any]:
         """Get current status information for a container.
@@ -269,7 +254,7 @@ class DockerModule:
             raise KeyError(f"Container {name} not found")
 
         try:
-            container = self.containers[name]
+            container, socket = self.containers[name]
             container.reload()  # Refresh container info
             
             return {

test/docker_module_test.py

@@ -5,175 +5,117 @@ from sia.docker_module import DockerModule
 class DockerModuleTest(unittest.TestCase):
     def test_initialization(self):
         """Test DockerModule initialization."""
-        docker_module = DockerModule()
+        with DockerModule() as docker_module:
-        self.assertIsInstance(docker_module, DockerModule)
+            self.assertIsInstance(docker_module, DockerModule)
 
     def test_start_container_short_lived(self):
         """Test starting a short-lived container."""
-        docker_module = DockerModule()
+        with DockerModule() as docker_module:
-        
+            image = "busybox:latest"
-        # Test parameters
+            timeout = 1000
-        image = "busybox:latest"
+            command = "echo"
-        timeout = 1000
+            arguments = ["Hello World"]
-        command = "echo"
+            output = docker_module.start_container(
-        arguments = ["Hello World"]
+                image=image,
-        
+                timeout=timeout,
-        output = docker_module.start_container(
+                command=command,
-            image=image,
+                arguments=arguments
-            timeout=timeout,
+            )
-            command=command,
+            self.assertEqual(output, "Hello World\n")
-            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()
+        with DockerModule() as docker_module:
-        
+            name = "test_start_container_long_running"
-        name = "test_start_container_long_running"
+            image = "busybox:latest"
-        image = "busybox:latest"
+            container_name = docker_module.start_container(
-        
+                image=image,
-        container_name = docker_module.start_container(
+                name=name
-            image=image,
+            )
-            name=name
+            self.assertEqual(container_name, name)
-        )
+            self.assertIn(name, docker_module.containers)
-        
+            docker_module.wait_container(container_name, 0)
-        # Verify container was started and stored
+            self.assertNotIn(name, docker_module.containers)
-        self.assertEqual(container_name, name)
-        self.assertIn(name, docker_module.containers)
 
     def test_start_container_with_volumes(self):
         """Test starting a container with volume mappings."""
-        docker_module = DockerModule()
+        with DockerModule() as docker_module:
-
+            arguments = ["sh", "-c", "echo 'Hello World' > /write_here/test_start_container_with_volumes"]
-        arguments = ["sh", "-c", "echo 'Hello World' > /write_here/test_start_container_with_volumes"]
+            volumes = {"/tmp": "/write_here"}
-        volumes = {"/tmp": "/write_here"}
+            docker_module.start_container(
-
+                image="busybox:latest",
-        output = docker_module.start_container(
+                timeout=1000,
-            image="busybox:latest",
+                arguments=arguments,
-            timeout=1000,
+                volumes=volumes
-            arguments=arguments,
+            )
-            volumes=volumes
+            with open("/tmp/test_start_container_with_volumes", "r") as f:
-        )
+                self.assertEqual(f.read(), "Hello World\n")
-        print(f"output: {output}", flush=True)
-        # Verify output was written to volume
-        with open("/tmp/test_start_container_with_volumes", "r") as f:
-            self.assertEqual(f.read(), "Hello World\n")
 
     def test_invalid_container_start(self):
         """Test error handling for invalid container start."""
-        docker_module = DockerModule()
+        with DockerModule() as docker_module:
-        
+            with self.assertRaises(ValueError):
-        # Test missing required parameters
+                docker_module.start_container(image="busybox:latest")
-        with self.assertRaises(ValueError):
-            docker_module.start_container(image="busybox:latest")
-
-    def test_stop_container(self):
-        """Test stopping a container."""
-        docker_module = DockerModule()
-        
-        # Start a container first
-        name = "test_stop_container"
-        docker_module.start_container(
-            image="busybox:latest",
-            name=name
-        )
-        
-        # Stop the container
-        docker_module.stop_container(name)
-        
-        # Verify container was stopped and removed
-        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()
+        with DockerModule() as docker_module:
+            arguments = [
+                "sh",
+                "-c",
+                "read input; echo \"$input\" | tr '[:lower:]' '[:upper:]'; echo \"$input\" | tr '[:upper:]' '[:lower:]' >&2"
+            ]
+            name = "test_container_io_operations"
+            container_name = docker_module.start_container(
+                image="busybox:latest",
+                name=name,
+                arguments=arguments
+            )
+            test_input = "Test Input\n"
+            docker_module.write_container_stdin(name, test_input)
+            time.sleep(1) # for processing
+            stdout = docker_module.read_container_stdout(name)
+            self.assertEqual(stdout, "TEST INPUT\n")
+            stderr = docker_module.read_container_stderr(name)
+            self.assertEqual(stderr, "test input\n")
+            docker_module.wait_container(container_name, 0)
 
-        arguments = [
+    def test_wait_container(self):
-            "sh",
+        """Test waiting for container completion."""
-            "-c",
+        with DockerModule() as docker_module:
-            "read input; echo \"$input\" | tr '[:lower:]' '[:upper:]'; echo \"$input\" | tr '[:upper:]' '[:lower:]' >&2"
+            name = "test_wait_container"
-        ]
+            docker_module.start_container(
+                image="busybox:latest",
+                name=name
+            )
+            exit_code, output = docker_module.wait_container(name, timeout=1000)
+            self.assertEqual(exit_code, -1)
+            self.assertEqual(output, "")
 
-        # Start a container
+    def test_get_container_status(self):
-        name = "test_container_io_operations"
+        """Test getting container status information."""
-        docker_module.start_container(
+        with DockerModule() as docker_module:
-            image="busybox:latest",
+            arguments = [
-            name=name,
+                "sh",
-            arguments=arguments
+                "-c",
-        )
+                "echo 'standard output'; echo 'standard error' >&2; sleep 3"
-        
+            ]
-        # Test stdin write
+            name = "test_get_container_status"
-        test_input = "Test Input\n"
+            docker_module.start_container(
-        docker_module.write_container_stdin(name, test_input)
+                image="busybox:latest",
-        time.sleep(1) # for processing
+                name=name,
-        
+                arguments=arguments
-        # Test stdout read
+            )
-        stdout = docker_module.read_container_stdout(name)
+            time.sleep(1)
-        print(f"stdout: {stdout}", flush=True)
+            status = docker_module.get_container_status(name)
-        self.assertEqual(stdout, "test input\n")
+            self.assertEqual(status['name'], name)
-        
+            self.assertEqual(status['status'], "running")
-        # Test stderr read
+            self.assertGreater(len(status['started_at']), 0)
-        stderr = docker_module.read_container_stderr(name)
+            self.assertEqual(status['exit_code'], 0)
-        print(f"stderr: {stderr}", flush=True)
+            self.assertEqual(status['error'], "")
-        self.assertEqual(stderr, "TEST INPUT\n")
+            self.assertEqual(status['stdout_size'], 16)
-
+            self.assertEqual(status['stderr_size'], 15)
-#    def test_wait_container(self):
+            docker_module.wait_container(name, 0)
-#        """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()