From 016187fa21def18eb74c229513135426527c8e23 Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Fri, 25 Oct 2024 14:37:56 +0200 Subject: [PATCH] Docker module basic functionality --- sia/docker_module.py | 179 +++++++++++++-------------- test/docker_module_test.py | 246 ++++++++++++++----------------------- 2 files changed, 176 insertions(+), 249 deletions(-) diff --git a/sia/docker_module.py b/sia/docker_module.py index fcdb6dd..5ea712a 100644 --- a/sia/docker_module.py +++ b/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.container_sockets: Dict[str, docker.APIClient.AttachSocket] = {} + self.containers: Dict[str, (Container, 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 - - try: - self.client.images.pull('busybox:latest') - # Create container with tty and open stdin - container = self.client.containers.create( - image, - command=cmd if cmd else None, - name=name, - detach=True, - volumes=volume_binds, - ports=port_bindings, - environment=environment, - tty=True, - stdin_open=True - ) - - # Attach socket before starting container - socket = container.attach_socket(params={ - 'stdin': 1, - 'stdout': 1, - 'stderr': 1, - 'stream': 1 - }) - - # Start the container - container.start() - - 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 + self.client.images.pull(image) + container = self.client.containers.create( + image, + command=cmd if cmd else None, + name=name, + detach=True, + volumes=volume_binds, + ports=port_bindings, + environment=environment, + tty=False, + stdin_open=True + ) + socket = container.attach_socket(params={ + 'stdin': 1, + 'stdout': 1, + 'stderr': 1, + 'stream': 1, + 'demux': True + }) + container.start() + if timeout >= 0: + # For short-lived containers + return self._wait_and_stop(container, socket, timeout)[1] + else: + # For long-running containers + self.containers[name] = (container, socket) + return name 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: - raise KeyError(f"Container {name} not found or socket not initialized") + if name not in self.containers: + 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 = self.containers[name] - 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 + container, socket = self.containers.pop(name) + return self._wait_and_stop(container, socket, timeout) 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 { diff --git a/test/docker_module_test.py b/test/docker_module_test.py index ea0862e..5841866 100644 --- a/test/docker_module_test.py +++ b/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() - self.assertIsInstance(docker_module, DockerModule) + with DockerModule() as docker_module: + self.assertIsInstance(docker_module, DockerModule) def test_start_container_short_lived(self): """Test starting a short-lived container.""" - docker_module = DockerModule() - - # Test parameters - image = "busybox: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") + with DockerModule() as docker_module: + image = "busybox:latest" + timeout = 1000 + command = "echo" + arguments = ["Hello World"] + output = docker_module.start_container( + image=image, + timeout=timeout, + command=command, + arguments=arguments + ) + self.assertEqual(output, "Hello World\n") def test_start_container_long_running(self): """Test starting a long-running container.""" - docker_module = DockerModule() - - name = "test_start_container_long_running" - image = "busybox: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) + with DockerModule() as docker_module: + name = "test_start_container_long_running" + image = "busybox:latest" + container_name = docker_module.start_container( + image=image, + name=name + ) + self.assertEqual(container_name, name) + self.assertIn(name, docker_module.containers) + docker_module.wait_container(container_name, 0) + self.assertNotIn(name, docker_module.containers) def test_start_container_with_volumes(self): """Test starting a container with volume mappings.""" - docker_module = DockerModule() - - arguments = ["sh", "-c", "echo 'Hello World' > /write_here/test_start_container_with_volumes"] - volumes = {"/tmp": "/write_here"} - - output = docker_module.start_container( - image="busybox:latest", - timeout=1000, - arguments=arguments, - volumes=volumes - ) - 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") + with DockerModule() as docker_module: + arguments = ["sh", "-c", "echo 'Hello World' > /write_here/test_start_container_with_volumes"] + volumes = {"/tmp": "/write_here"} + docker_module.start_container( + image="busybox:latest", + timeout=1000, + arguments=arguments, + volumes=volumes + ) + 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() - - # Test missing required parameters - 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") + with DockerModule() as docker_module: + with self.assertRaises(ValueError): + docker_module.start_container(image="busybox:latest") 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 = [ - "sh", - "-c", - "read input; echo \"$input\" | tr '[:lower:]' '[:upper:]'; echo \"$input\" | tr '[:upper:]' '[:lower:]' >&2" - ] - - # Start a container - name = "test_container_io_operations" - docker_module.start_container( - image="busybox:latest", - name=name, - arguments=arguments - ) - - # Test stdin write - test_input = "Test Input\n" - docker_module.write_container_stdin(name, test_input) - time.sleep(1) # for processing - - # Test stdout read - stdout = docker_module.read_container_stdout(name) - print(f"stdout: {stdout}", flush=True) - self.assertEqual(stdout, "test input\n") - - # Test stderr read - stderr = docker_module.read_container_stderr(name) - print(f"stderr: {stderr}", flush=True) - self.assertEqual(stderr, "TEST INPUT\n") + def test_wait_container(self): + """Test waiting for container completion.""" + with DockerModule() as docker_module: + 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, "") -# 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() + def test_get_container_status(self): + """Test getting container status information.""" + with DockerModule() as docker_module: + arguments = [ + "sh", + "-c", + "echo 'standard output'; echo 'standard error' >&2; sleep 3" + ] + name = "test_get_container_status" + docker_module.start_container( + image="busybox:latest", + name=name, + arguments=arguments + ) + time.sleep(1) + status = docker_module.get_container_status(name) + self.assertEqual(status['name'], name) + self.assertEqual(status['status'], "running") + self.assertGreater(len(status['started_at']), 0) + self.assertEqual(status['exit_code'], 0) + self.assertEqual(status['error'], "") + self.assertEqual(status['stdout_size'], 16) + self.assertEqual(status['stderr_size'], 15) + docker_module.wait_container(name, 0) if __name__ == '__main__': unittest.main()