From 8ff8f1b16141e03d9160e8c43d28b9a6a2de5b4b Mon Sep 17 00:00:00 2001 From: Hamdi Sakly Date: Tue, 14 Apr 2026 15:00:38 +0100 Subject: [PATCH 1/7] feat : input and output files --- .../code_executors/container_code_executor.py | 401 +++++++++++------- .../test_container_code_executor.py | 257 +++++++++++ 2 files changed, 511 insertions(+), 147 deletions(-) create mode 100644 tests/unittests/code_executors/test_container_code_executor.py diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index d6a78d4d26..96f25825a6 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -15,8 +15,10 @@ from __future__ import annotations import atexit +import io import logging import os +import tarfile from typing import Optional import docker @@ -29,172 +31,277 @@ from .base_code_executor import BaseCodeExecutor from .code_execution_utils import CodeExecutionInput from .code_execution_utils import CodeExecutionResult +from .code_execution_utils import File -logger = logging.getLogger('google_adk.' + __name__) -DEFAULT_IMAGE_TAG = 'adk-code-executor:latest' +logger = logging.getLogger("google_adk." + __name__) +DEFAULT_IMAGE_TAG = "adk-code-executor:latest" class ContainerCodeExecutor(BaseCodeExecutor): - """A code executor that uses a custom container to execute code. - - Attributes: - base_url: Optional. The base url of the user hosted Docker client. - image: The tag of the predefined image or custom image to run on the - container. Either docker_path or image must be set. - docker_path: The path to the directory containing the Dockerfile. If set, - build the image from the dockerfile path instead of using the predefined - image. Either docker_path or image must be set. - """ + """A code executor that uses a custom container to execute code. - base_url: Optional[str] = None - """ + Attributes: + base_url: Optional. The base url of the user hosted Docker client. + image: The tag of the predefined image or custom image to run on the + container. Either docker_path or image must be set. + docker_path: The path to the directory containing the Dockerfile. If set, + build the image from the dockerfile path instead of using the predefined + image. Either docker_path or image must be set. + input_dir: The directory in the container where input files will be placed. + output_dir: The directory in the container where output files will be + retrieved from. + """ + + base_url: Optional[str] = None + """ Optional. The base url of the user hosted Docker client. """ - image: str = None - """ + image: str = None + """ The tag of the predefined image or custom image to run on the container. Either docker_path or image must be set. """ - docker_path: str = None - """ + docker_path: str = None + """ The path to the directory containing the Dockerfile. If set, build the image from the dockerfile path instead of using the predefined image. Either docker_path or image must be set. """ - # Overrides the BaseCodeExecutor attribute: this executor cannot be stateful. - stateful: bool = Field(default=False, frozen=True, exclude=True) + input_dir: str = "/tmp/inputs" + """ + The directory in the container where input files will be placed. + """ - # Overrides the BaseCodeExecutor attribute: this executor cannot - # optimize_data_file. - optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) + output_dir: str = "/tmp/outputs" + """ + The directory in the container where output files will be retrieved from. + """ - _client: DockerClient = None - _container: Container = None + stateful: bool = Field(default=False, frozen=True, exclude=True) - def __init__( - self, - base_url: Optional[str] = None, - image: Optional[str] = None, - docker_path: Optional[str] = None, - **data, - ): - """Initializes the ContainerCodeExecutor. + optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) - Args: - base_url: Optional. The base url of the user hosted Docker client. - image: The tag of the predefined image or custom image to run on the - container. Either docker_path or image must be set. - docker_path: The path to the directory containing the Dockerfile. If set, - build the image from the dockerfile path instead of using the predefined - image. Either docker_path or image must be set. - **data: The data to initialize the ContainerCodeExecutor. - """ - if not image and not docker_path: - raise ValueError( - 'Either image or docker_path must be set for ContainerCodeExecutor.' - ) - if 'stateful' in data and data['stateful']: - raise ValueError('Cannot set `stateful=True` in ContainerCodeExecutor.') - if 'optimize_data_file' in data and data['optimize_data_file']: - raise ValueError( - 'Cannot set `optimize_data_file=True` in ContainerCodeExecutor.' - ) - - super().__init__(**data) - self.base_url = base_url - self.image = image if image else DEFAULT_IMAGE_TAG - self.docker_path = os.path.abspath(docker_path) if docker_path else None - - self._client = ( - docker.from_env() - if not self.base_url - else docker.DockerClient(base_url=self.base_url) - ) - # Initialize the container. - self.__init_container() - - # Close the container when the on exit. - atexit.register(self.__cleanup_container) - - @override - def execute_code( - self, - invocation_context: InvocationContext, - code_execution_input: CodeExecutionInput, - ) -> CodeExecutionResult: - output = '' - error = '' - exec_result = self._container.exec_run( - ['python3', '-c', code_execution_input.code], - demux=True, - ) - logger.debug('Executed code:\n```\n%s\n```', code_execution_input.code) - - if exec_result.output and exec_result.output[0]: - output = exec_result.output[0].decode('utf-8') - if ( - exec_result.output - and len(exec_result.output) > 1 - and exec_result.output[1] + _client: DockerClient = None + _container: Container = None + + def __init__( + self, + base_url: Optional[str] = None, + image: Optional[str] = None, + docker_path: Optional[str] = None, + input_dir: Optional[str] = None, + output_dir: Optional[str] = None, + **data, ): - error = exec_result.output[1].decode('utf-8') - - # Collect the final result. - return CodeExecutionResult( - stdout=output, - stderr=error, - output_files=[], - ) - - def _build_docker_image(self): - """Builds the Docker image.""" - if not self.docker_path: - raise ValueError('Docker path is not set.') - if not os.path.exists(self.docker_path): - raise FileNotFoundError(f'Invalid Docker path: {self.docker_path}') - - logger.info('Building Docker image...') - self._client.images.build( - path=self.docker_path, - tag=self.image, - rm=True, - ) - logger.info('Docker image: %s built.', self.image) - - def _verify_python_installation(self): - """Verifies the container has python3 installed.""" - exec_result = self._container.exec_run(['which', 'python3']) - if exec_result.exit_code != 0: - raise ValueError('python3 is not installed in the container.') - - def __init_container(self): - """Initializes the container.""" - if not self._client: - raise RuntimeError('Docker client is not initialized.') - - if self.docker_path: - self._build_docker_image() - - logger.info('Starting container for ContainerCodeExecutor...') - self._container = self._client.containers.run( - image=self.image, - detach=True, - tty=True, - ) - logger.info('Container %s started.', self._container.id) - - # Verify the container is able to run python3. - self._verify_python_installation() - - def __cleanup_container(self): - """Closes the container on exit.""" - if not self._container: - return - - logger.info('[Cleanup] Stopping the container...') - self._container.stop() - self._container.remove() - logger.info('Container %s stopped and removed.', self._container.id) + """Initializes the ContainerCodeExecutor. + + Args: + base_url: Optional. The base url of the user hosted Docker client. + image: The tag of the predefined image or custom image to run on the + container. Either docker_path or image must be set. + docker_path: The path to the directory containing the Dockerfile. If set, + build the image from the dockerfile path instead of using the predefined + image. Either docker_path or image must be set. + input_dir: The directory in the container where input files will be placed. + Defaults to '/tmp/inputs'. + output_dir: The directory in the container where output files will be + retrieved from. Defaults to '/tmp/outputs'. + **data: The data to initialize the ContainerCodeExecutor. + """ + if not image and not docker_path: + raise ValueError( + "Either image or docker_path must be set for ContainerCodeExecutor." + ) + if "stateful" in data and data["stateful"]: + raise ValueError("Cannot set `stateful=True` in ContainerCodeExecutor.") + if "optimize_data_file" in data and data["optimize_data_file"]: + raise ValueError( + "Cannot set `optimize_data_file=True` in ContainerCodeExecutor." + ) + + super().__init__(**data) + self.base_url = base_url + self.image = image if image else DEFAULT_IMAGE_TAG + self.docker_path = os.path.abspath(docker_path) if docker_path else None + self.input_dir = input_dir if input_dir else "/tmp/inputs" + self.output_dir = output_dir if output_dir else "/tmp/outputs" + + self._client = ( + docker.from_env() + if not self.base_url + else docker.DockerClient(base_url=self.base_url) + ) + self.__init_container() + + atexit.register(self.__cleanup_container) + + @override + def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + if code_execution_input.input_files: + self._put_input_files(code_execution_input.input_files) + + self._create_output_directory() + + output = "" + error = "" + exec_result = self._container.exec_run( + ["python3", "-c", code_execution_input.code], + demux=True, + ) + logger.debug("Executed code:\n```\n%s\n```", code_execution_input.code) + + if exec_result.output and exec_result.output[0]: + output = exec_result.output[0].decode("utf-8") + if exec_result.output and len(exec_result.output) > 1 and exec_result.output[1]: + error = exec_result.output[1].decode("utf-8") + + output_files = self._get_output_files() + + return CodeExecutionResult( + stdout=output, + stderr=error, + output_files=output_files, + ) + + def _build_docker_image(self): + """Builds the Docker image.""" + if not self.docker_path: + raise ValueError("Docker path is not set.") + if not os.path.exists(self.docker_path): + raise FileNotFoundError(f"Invalid Docker path: {self.docker_path}") + + logger.info("Building Docker image...") + self._client.images.build( + path=self.docker_path, + tag=self.image, + rm=True, + ) + logger.info("Docker image: %s built.", self.image) + + def _verify_python_installation(self): + """Verifies the container has python3 installed.""" + exec_result = self._container.exec_run(["which", "python3"]) + if exec_result.exit_code != 0: + raise ValueError("python3 is not installed in the container.") + + def _put_input_files(self, input_files: list[File]) -> None: + """Puts input files into the container using put_archive. + + Args: + input_files: The list of input files to copy into the container. + """ + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + for file in input_files: + content = file.content + if isinstance(content, str): + content = content.encode("utf-8") + tarinfo = tarfile.TarInfo(name=file.name) + tarinfo.size = len(content) + tar.addfile(tarinfo, io.BytesIO(content)) + + tar_buffer.seek(0) + self._container.put_archive( + self.input_dir, + tar_buffer.read(), + ) + logger.debug("Copied %d input files to %s", len(input_files), self.input_dir) + + def _create_output_directory(self) -> None: + """Creates the output directory in the container if it doesn't exist.""" + exec_result = self._container.exec_run( + ["mkdir", "-p", self.output_dir], + ) + if exec_result.exit_code != 0: + logger.warning( + "Failed to create output directory %s: %s", + self.output_dir, + exec_result.output, + ) + + def _get_output_files(self) -> list[File]: + """Gets output files from the container. + + Returns: + The list of output files retrieved from the container. + """ + try: + tar_bytes, stat = self._container.get_archive(self.output_dir) + except docker.errors.APIError as e: + if e.response.status_code == 404: + logger.debug("No output files found at %s", self.output_dir) + return [] + raise + + tar_buffer = io.BytesIO(tar_bytes) + output_files = [] + + with tarfile.open(fileobj=tar_buffer, mode="r") as tar: + for member in tar.getmembers(): + if member.isfile(): + file_obj = tar.extractfile(member) + if file_obj: + content = file_obj.read() + file_name = os.path.basename(member.name) + if file_name: + output_files.append( + File( + name=file_name, + content=content, + mime_type=self._guess_mime_type(file_name), + ) + ) + + logger.debug( + "Retrieved %d output files from %s", len(output_files), self.output_dir + ) + return output_files + + def _guess_mime_type(self, file_name: str) -> str: + """Guesses the MIME type based on the file extension. + + Args: + file_name: The name of the file. + + Returns: + The guessed MIME type, or 'application/octet-stream' if unknown. + """ + import mimetypes + + mime_type, _ = mimetypes.guess_type(file_name) + return mime_type if mime_type else "application/octet-stream" + + def __init_container(self): + """Initializes the container.""" + if not self._client: + raise RuntimeError("Docker client is not initialized.") + + if self.docker_path: + self._build_docker_image() + + logger.info("Starting container for ContainerCodeExecutor...") + self._container = self._client.containers.run( + image=self.image, + detach=True, + tty=True, + ) + logger.info("Container %s started.", self._container.id) + + self._verify_python_installation() + + def __cleanup_container(self): + """Closes the container on exit.""" + if not self._container: + return + + logger.info("[Cleanup] Stopping the container...") + self._container.stop() + self._container.remove() + logger.info("Container %s stopped and removed.", self._container.id) diff --git a/tests/unittests/code_executors/test_container_code_executor.py b/tests/unittests/code_executors/test_container_code_executor.py new file mode 100644 index 0000000000..b49c186134 --- /dev/null +++ b/tests/unittests/code_executors/test_container_code_executor.py @@ -0,0 +1,257 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ContainerCodeExecutor.""" + +import io +import tarfile +from unittest import mock + +from google.adk.agents.invocation_context import InvocationContext +from google.adk.code_executors.code_execution_utils import CodeExecutionInput +from google.adk.code_executors.code_execution_utils import CodeExecutionResult +from google.adk.code_executors.code_execution_utils import File +from google.adk.code_executors.container_code_executor import ContainerCodeExecutor +import pytest + + +@pytest.fixture +def mock_container(): + container = mock.MagicMock() + container.id = "test-container-id" + container.exec_run.return_value = mock.MagicMock( + output=(b"Hello World", b""), + exit_code=0, + ) + return container + + +@pytest.fixture +def mock_docker_client(mock_container): + client = mock.MagicMock() + client.containers.run.return_value = mock_container + return client + + +class TestContainerCodeExecutorInit: + + def test_init_with_image(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + assert executor.image == "test-image:latest" + assert executor.input_dir == "/tmp/inputs" + assert executor.output_dir == "/tmp/outputs" + + def test_init_with_custom_dirs(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor( + image="test-image:latest", + input_dir="/custom/inputs", + output_dir="/custom/outputs", + ) + assert executor.input_dir == "/custom/inputs" + assert executor.output_dir == "/custom/outputs" + + def test_init_requires_image_or_docker_path(self): + with pytest.raises( + ValueError, match="Either image or docker_path must be set" + ): + ContainerCodeExecutor() + + def test_init_rejects_stateful(self): + with pytest.raises(ValueError, match="Cannot set `stateful=True`"): + ContainerCodeExecutor(image="test", stateful=True) + + def test_init_rejects_optimize_data_file(self): + with pytest.raises( + ValueError, match="Cannot set `optimize_data_file=True`" + ): + ContainerCodeExecutor(image="test", optimize_data_file=True) + + +class TestExecuteCode: + + def test_execute_code_basic(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("Hello World")') + + result = executor.execute_code(context, code_input) + + assert result.stdout == "Hello World" + assert result.stderr == "" + assert result.output_files == [] + + def test_execute_code_with_error(self, mock_docker_client, mock_container): + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b"Some error"), + exit_code=1, + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='raise Error("test")') + + result = executor.execute_code(context, code_input) + + assert result.stderr == "Some error" + + def test_execute_code_with_input_files( + self, mock_docker_client, mock_container + ): + mock_container.put_archive = mock.MagicMock() + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + mock_container.get_archive = mock.MagicMock( + side_effect=mock.MagicMock(response=mock.MagicMock(status_code=404)), + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput( + code='print("test")', + input_files=[File(name="test.txt", content="test content")], + ) + + result = executor.execute_code(context, code_input) + + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" + + def test_execute_code_with_output_files( + self, mock_docker_client, mock_container + ): + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar.addfile( + tarfile.TarInfo(name="output.txt"), + io.BytesIO(b"output content"), + ) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() + + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("test")') + + result = executor.execute_code(context, code_input) + + assert len(result.output_files) == 1 + assert result.output_files[0].name == "output.txt" + + +class TestPutInputFiles: + + def test_put_archive_called(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [ + File(name="file1.txt", content="content1"), + File(name="file2.txt", content="content2"), + ] + + executor._put_input_files(input_files) + + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" + + def test_handles_string_content(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [File(name="test.txt", content="string content")] + + executor._put_input_files(input_files) + + mock_container.put_archive.assert_called_once() + + +class TestGetOutputFiles: + + def test_no_output_files(self, mock_docker_client, mock_container): + import docker + + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + output_files = executor._get_output_files() + + assert output_files == [] + + def test_extracts_files_from_archive( + self, mock_docker_client, mock_container + ): + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar.addfile( + tarfile.TarInfo(name="output.txt"), + io.BytesIO(b"output content"), + ) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() + + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + output_files = executor._get_output_files() + + assert len(output_files) == 1 + assert output_files[0].name == "output.txt" + assert output_files[0].content == b"output content" + + +class TestMimeTypeGuessing: + + def test_guess_txt(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + mime_type = executor._guess_mime_type("test.txt") + + assert mime_type == "text/plain" + + def test_guess_csv(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + mime_type = executor._guess_mime_type("data.csv") + + assert mime_type == "text/csv" + + def test_default_for_unknown(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + mime_type = executor._guess_mime_type("unknown.xyz") + + assert mime_type == "application/octet-stream" From 5591087dd053933775d1e06198da339d13fc0718 Mon Sep 17 00:00:00 2001 From: Hamdi Sakly Date: Tue, 14 Apr 2026 15:00:38 +0100 Subject: [PATCH 2/7] feat : input and output files --- .../code_executors/container_code_executor.py | 401 +++++++++++------- .../test_container_code_executor.py | 257 +++++++++++ 2 files changed, 511 insertions(+), 147 deletions(-) create mode 100644 tests/unittests/code_executors/test_container_code_executor.py diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index d6a78d4d26..96f25825a6 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -15,8 +15,10 @@ from __future__ import annotations import atexit +import io import logging import os +import tarfile from typing import Optional import docker @@ -29,172 +31,277 @@ from .base_code_executor import BaseCodeExecutor from .code_execution_utils import CodeExecutionInput from .code_execution_utils import CodeExecutionResult +from .code_execution_utils import File -logger = logging.getLogger('google_adk.' + __name__) -DEFAULT_IMAGE_TAG = 'adk-code-executor:latest' +logger = logging.getLogger("google_adk." + __name__) +DEFAULT_IMAGE_TAG = "adk-code-executor:latest" class ContainerCodeExecutor(BaseCodeExecutor): - """A code executor that uses a custom container to execute code. - - Attributes: - base_url: Optional. The base url of the user hosted Docker client. - image: The tag of the predefined image or custom image to run on the - container. Either docker_path or image must be set. - docker_path: The path to the directory containing the Dockerfile. If set, - build the image from the dockerfile path instead of using the predefined - image. Either docker_path or image must be set. - """ + """A code executor that uses a custom container to execute code. - base_url: Optional[str] = None - """ + Attributes: + base_url: Optional. The base url of the user hosted Docker client. + image: The tag of the predefined image or custom image to run on the + container. Either docker_path or image must be set. + docker_path: The path to the directory containing the Dockerfile. If set, + build the image from the dockerfile path instead of using the predefined + image. Either docker_path or image must be set. + input_dir: The directory in the container where input files will be placed. + output_dir: The directory in the container where output files will be + retrieved from. + """ + + base_url: Optional[str] = None + """ Optional. The base url of the user hosted Docker client. """ - image: str = None - """ + image: str = None + """ The tag of the predefined image or custom image to run on the container. Either docker_path or image must be set. """ - docker_path: str = None - """ + docker_path: str = None + """ The path to the directory containing the Dockerfile. If set, build the image from the dockerfile path instead of using the predefined image. Either docker_path or image must be set. """ - # Overrides the BaseCodeExecutor attribute: this executor cannot be stateful. - stateful: bool = Field(default=False, frozen=True, exclude=True) + input_dir: str = "/tmp/inputs" + """ + The directory in the container where input files will be placed. + """ - # Overrides the BaseCodeExecutor attribute: this executor cannot - # optimize_data_file. - optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) + output_dir: str = "/tmp/outputs" + """ + The directory in the container where output files will be retrieved from. + """ - _client: DockerClient = None - _container: Container = None + stateful: bool = Field(default=False, frozen=True, exclude=True) - def __init__( - self, - base_url: Optional[str] = None, - image: Optional[str] = None, - docker_path: Optional[str] = None, - **data, - ): - """Initializes the ContainerCodeExecutor. + optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) - Args: - base_url: Optional. The base url of the user hosted Docker client. - image: The tag of the predefined image or custom image to run on the - container. Either docker_path or image must be set. - docker_path: The path to the directory containing the Dockerfile. If set, - build the image from the dockerfile path instead of using the predefined - image. Either docker_path or image must be set. - **data: The data to initialize the ContainerCodeExecutor. - """ - if not image and not docker_path: - raise ValueError( - 'Either image or docker_path must be set for ContainerCodeExecutor.' - ) - if 'stateful' in data and data['stateful']: - raise ValueError('Cannot set `stateful=True` in ContainerCodeExecutor.') - if 'optimize_data_file' in data and data['optimize_data_file']: - raise ValueError( - 'Cannot set `optimize_data_file=True` in ContainerCodeExecutor.' - ) - - super().__init__(**data) - self.base_url = base_url - self.image = image if image else DEFAULT_IMAGE_TAG - self.docker_path = os.path.abspath(docker_path) if docker_path else None - - self._client = ( - docker.from_env() - if not self.base_url - else docker.DockerClient(base_url=self.base_url) - ) - # Initialize the container. - self.__init_container() - - # Close the container when the on exit. - atexit.register(self.__cleanup_container) - - @override - def execute_code( - self, - invocation_context: InvocationContext, - code_execution_input: CodeExecutionInput, - ) -> CodeExecutionResult: - output = '' - error = '' - exec_result = self._container.exec_run( - ['python3', '-c', code_execution_input.code], - demux=True, - ) - logger.debug('Executed code:\n```\n%s\n```', code_execution_input.code) - - if exec_result.output and exec_result.output[0]: - output = exec_result.output[0].decode('utf-8') - if ( - exec_result.output - and len(exec_result.output) > 1 - and exec_result.output[1] + _client: DockerClient = None + _container: Container = None + + def __init__( + self, + base_url: Optional[str] = None, + image: Optional[str] = None, + docker_path: Optional[str] = None, + input_dir: Optional[str] = None, + output_dir: Optional[str] = None, + **data, ): - error = exec_result.output[1].decode('utf-8') - - # Collect the final result. - return CodeExecutionResult( - stdout=output, - stderr=error, - output_files=[], - ) - - def _build_docker_image(self): - """Builds the Docker image.""" - if not self.docker_path: - raise ValueError('Docker path is not set.') - if not os.path.exists(self.docker_path): - raise FileNotFoundError(f'Invalid Docker path: {self.docker_path}') - - logger.info('Building Docker image...') - self._client.images.build( - path=self.docker_path, - tag=self.image, - rm=True, - ) - logger.info('Docker image: %s built.', self.image) - - def _verify_python_installation(self): - """Verifies the container has python3 installed.""" - exec_result = self._container.exec_run(['which', 'python3']) - if exec_result.exit_code != 0: - raise ValueError('python3 is not installed in the container.') - - def __init_container(self): - """Initializes the container.""" - if not self._client: - raise RuntimeError('Docker client is not initialized.') - - if self.docker_path: - self._build_docker_image() - - logger.info('Starting container for ContainerCodeExecutor...') - self._container = self._client.containers.run( - image=self.image, - detach=True, - tty=True, - ) - logger.info('Container %s started.', self._container.id) - - # Verify the container is able to run python3. - self._verify_python_installation() - - def __cleanup_container(self): - """Closes the container on exit.""" - if not self._container: - return - - logger.info('[Cleanup] Stopping the container...') - self._container.stop() - self._container.remove() - logger.info('Container %s stopped and removed.', self._container.id) + """Initializes the ContainerCodeExecutor. + + Args: + base_url: Optional. The base url of the user hosted Docker client. + image: The tag of the predefined image or custom image to run on the + container. Either docker_path or image must be set. + docker_path: The path to the directory containing the Dockerfile. If set, + build the image from the dockerfile path instead of using the predefined + image. Either docker_path or image must be set. + input_dir: The directory in the container where input files will be placed. + Defaults to '/tmp/inputs'. + output_dir: The directory in the container where output files will be + retrieved from. Defaults to '/tmp/outputs'. + **data: The data to initialize the ContainerCodeExecutor. + """ + if not image and not docker_path: + raise ValueError( + "Either image or docker_path must be set for ContainerCodeExecutor." + ) + if "stateful" in data and data["stateful"]: + raise ValueError("Cannot set `stateful=True` in ContainerCodeExecutor.") + if "optimize_data_file" in data and data["optimize_data_file"]: + raise ValueError( + "Cannot set `optimize_data_file=True` in ContainerCodeExecutor." + ) + + super().__init__(**data) + self.base_url = base_url + self.image = image if image else DEFAULT_IMAGE_TAG + self.docker_path = os.path.abspath(docker_path) if docker_path else None + self.input_dir = input_dir if input_dir else "/tmp/inputs" + self.output_dir = output_dir if output_dir else "/tmp/outputs" + + self._client = ( + docker.from_env() + if not self.base_url + else docker.DockerClient(base_url=self.base_url) + ) + self.__init_container() + + atexit.register(self.__cleanup_container) + + @override + def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + if code_execution_input.input_files: + self._put_input_files(code_execution_input.input_files) + + self._create_output_directory() + + output = "" + error = "" + exec_result = self._container.exec_run( + ["python3", "-c", code_execution_input.code], + demux=True, + ) + logger.debug("Executed code:\n```\n%s\n```", code_execution_input.code) + + if exec_result.output and exec_result.output[0]: + output = exec_result.output[0].decode("utf-8") + if exec_result.output and len(exec_result.output) > 1 and exec_result.output[1]: + error = exec_result.output[1].decode("utf-8") + + output_files = self._get_output_files() + + return CodeExecutionResult( + stdout=output, + stderr=error, + output_files=output_files, + ) + + def _build_docker_image(self): + """Builds the Docker image.""" + if not self.docker_path: + raise ValueError("Docker path is not set.") + if not os.path.exists(self.docker_path): + raise FileNotFoundError(f"Invalid Docker path: {self.docker_path}") + + logger.info("Building Docker image...") + self._client.images.build( + path=self.docker_path, + tag=self.image, + rm=True, + ) + logger.info("Docker image: %s built.", self.image) + + def _verify_python_installation(self): + """Verifies the container has python3 installed.""" + exec_result = self._container.exec_run(["which", "python3"]) + if exec_result.exit_code != 0: + raise ValueError("python3 is not installed in the container.") + + def _put_input_files(self, input_files: list[File]) -> None: + """Puts input files into the container using put_archive. + + Args: + input_files: The list of input files to copy into the container. + """ + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + for file in input_files: + content = file.content + if isinstance(content, str): + content = content.encode("utf-8") + tarinfo = tarfile.TarInfo(name=file.name) + tarinfo.size = len(content) + tar.addfile(tarinfo, io.BytesIO(content)) + + tar_buffer.seek(0) + self._container.put_archive( + self.input_dir, + tar_buffer.read(), + ) + logger.debug("Copied %d input files to %s", len(input_files), self.input_dir) + + def _create_output_directory(self) -> None: + """Creates the output directory in the container if it doesn't exist.""" + exec_result = self._container.exec_run( + ["mkdir", "-p", self.output_dir], + ) + if exec_result.exit_code != 0: + logger.warning( + "Failed to create output directory %s: %s", + self.output_dir, + exec_result.output, + ) + + def _get_output_files(self) -> list[File]: + """Gets output files from the container. + + Returns: + The list of output files retrieved from the container. + """ + try: + tar_bytes, stat = self._container.get_archive(self.output_dir) + except docker.errors.APIError as e: + if e.response.status_code == 404: + logger.debug("No output files found at %s", self.output_dir) + return [] + raise + + tar_buffer = io.BytesIO(tar_bytes) + output_files = [] + + with tarfile.open(fileobj=tar_buffer, mode="r") as tar: + for member in tar.getmembers(): + if member.isfile(): + file_obj = tar.extractfile(member) + if file_obj: + content = file_obj.read() + file_name = os.path.basename(member.name) + if file_name: + output_files.append( + File( + name=file_name, + content=content, + mime_type=self._guess_mime_type(file_name), + ) + ) + + logger.debug( + "Retrieved %d output files from %s", len(output_files), self.output_dir + ) + return output_files + + def _guess_mime_type(self, file_name: str) -> str: + """Guesses the MIME type based on the file extension. + + Args: + file_name: The name of the file. + + Returns: + The guessed MIME type, or 'application/octet-stream' if unknown. + """ + import mimetypes + + mime_type, _ = mimetypes.guess_type(file_name) + return mime_type if mime_type else "application/octet-stream" + + def __init_container(self): + """Initializes the container.""" + if not self._client: + raise RuntimeError("Docker client is not initialized.") + + if self.docker_path: + self._build_docker_image() + + logger.info("Starting container for ContainerCodeExecutor...") + self._container = self._client.containers.run( + image=self.image, + detach=True, + tty=True, + ) + logger.info("Container %s started.", self._container.id) + + self._verify_python_installation() + + def __cleanup_container(self): + """Closes the container on exit.""" + if not self._container: + return + + logger.info("[Cleanup] Stopping the container...") + self._container.stop() + self._container.remove() + logger.info("Container %s stopped and removed.", self._container.id) diff --git a/tests/unittests/code_executors/test_container_code_executor.py b/tests/unittests/code_executors/test_container_code_executor.py new file mode 100644 index 0000000000..b49c186134 --- /dev/null +++ b/tests/unittests/code_executors/test_container_code_executor.py @@ -0,0 +1,257 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ContainerCodeExecutor.""" + +import io +import tarfile +from unittest import mock + +from google.adk.agents.invocation_context import InvocationContext +from google.adk.code_executors.code_execution_utils import CodeExecutionInput +from google.adk.code_executors.code_execution_utils import CodeExecutionResult +from google.adk.code_executors.code_execution_utils import File +from google.adk.code_executors.container_code_executor import ContainerCodeExecutor +import pytest + + +@pytest.fixture +def mock_container(): + container = mock.MagicMock() + container.id = "test-container-id" + container.exec_run.return_value = mock.MagicMock( + output=(b"Hello World", b""), + exit_code=0, + ) + return container + + +@pytest.fixture +def mock_docker_client(mock_container): + client = mock.MagicMock() + client.containers.run.return_value = mock_container + return client + + +class TestContainerCodeExecutorInit: + + def test_init_with_image(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + assert executor.image == "test-image:latest" + assert executor.input_dir == "/tmp/inputs" + assert executor.output_dir == "/tmp/outputs" + + def test_init_with_custom_dirs(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor( + image="test-image:latest", + input_dir="/custom/inputs", + output_dir="/custom/outputs", + ) + assert executor.input_dir == "/custom/inputs" + assert executor.output_dir == "/custom/outputs" + + def test_init_requires_image_or_docker_path(self): + with pytest.raises( + ValueError, match="Either image or docker_path must be set" + ): + ContainerCodeExecutor() + + def test_init_rejects_stateful(self): + with pytest.raises(ValueError, match="Cannot set `stateful=True`"): + ContainerCodeExecutor(image="test", stateful=True) + + def test_init_rejects_optimize_data_file(self): + with pytest.raises( + ValueError, match="Cannot set `optimize_data_file=True`" + ): + ContainerCodeExecutor(image="test", optimize_data_file=True) + + +class TestExecuteCode: + + def test_execute_code_basic(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("Hello World")') + + result = executor.execute_code(context, code_input) + + assert result.stdout == "Hello World" + assert result.stderr == "" + assert result.output_files == [] + + def test_execute_code_with_error(self, mock_docker_client, mock_container): + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b"Some error"), + exit_code=1, + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='raise Error("test")') + + result = executor.execute_code(context, code_input) + + assert result.stderr == "Some error" + + def test_execute_code_with_input_files( + self, mock_docker_client, mock_container + ): + mock_container.put_archive = mock.MagicMock() + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + mock_container.get_archive = mock.MagicMock( + side_effect=mock.MagicMock(response=mock.MagicMock(status_code=404)), + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput( + code='print("test")', + input_files=[File(name="test.txt", content="test content")], + ) + + result = executor.execute_code(context, code_input) + + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" + + def test_execute_code_with_output_files( + self, mock_docker_client, mock_container + ): + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar.addfile( + tarfile.TarInfo(name="output.txt"), + io.BytesIO(b"output content"), + ) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() + + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("test")') + + result = executor.execute_code(context, code_input) + + assert len(result.output_files) == 1 + assert result.output_files[0].name == "output.txt" + + +class TestPutInputFiles: + + def test_put_archive_called(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [ + File(name="file1.txt", content="content1"), + File(name="file2.txt", content="content2"), + ] + + executor._put_input_files(input_files) + + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" + + def test_handles_string_content(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [File(name="test.txt", content="string content")] + + executor._put_input_files(input_files) + + mock_container.put_archive.assert_called_once() + + +class TestGetOutputFiles: + + def test_no_output_files(self, mock_docker_client, mock_container): + import docker + + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + output_files = executor._get_output_files() + + assert output_files == [] + + def test_extracts_files_from_archive( + self, mock_docker_client, mock_container + ): + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar.addfile( + tarfile.TarInfo(name="output.txt"), + io.BytesIO(b"output content"), + ) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() + + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + output_files = executor._get_output_files() + + assert len(output_files) == 1 + assert output_files[0].name == "output.txt" + assert output_files[0].content == b"output content" + + +class TestMimeTypeGuessing: + + def test_guess_txt(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + mime_type = executor._guess_mime_type("test.txt") + + assert mime_type == "text/plain" + + def test_guess_csv(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + mime_type = executor._guess_mime_type("data.csv") + + assert mime_type == "text/csv" + + def test_default_for_unknown(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + mime_type = executor._guess_mime_type("unknown.xyz") + + assert mime_type == "application/octet-stream" From bf7876edf4f5c08cb9b10a26d432bb6359ff9347 Mon Sep 17 00:00:00 2001 From: Hamdi Sakly Date: Wed, 15 Apr 2026 12:56:48 +0100 Subject: [PATCH 3/7] to rerun the adk script --- src/google/adk/code_executors/container_code_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index 96f25825a6..6f4a1b493e 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -67,7 +67,7 @@ class ContainerCodeExecutor(BaseCodeExecutor): """ The path to the directory containing the Dockerfile. If set, build the image from the dockerfile path instead of using the - predefined image. Either docker_path or image must be set. + predefined image. Either docker_path or image must be set . """ input_dir: str = "/tmp/inputs" From 0c520421483ccec3215590a0318b3080fcf660ac Mon Sep 17 00:00:00 2001 From: Hamdi Sakly Date: Thu, 16 Apr 2026 22:40:34 +0100 Subject: [PATCH 4/7] fix formatting --- .../code_executors/container_code_executor.py | 494 +++++++++--------- 1 file changed, 250 insertions(+), 244 deletions(-) diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index 6f4a1b493e..d301b32414 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -38,270 +38,276 @@ class ContainerCodeExecutor(BaseCodeExecutor): - """A code executor that uses a custom container to execute code. - - Attributes: - base_url: Optional. The base url of the user hosted Docker client. - image: The tag of the predefined image or custom image to run on the - container. Either docker_path or image must be set. - docker_path: The path to the directory containing the Dockerfile. If set, - build the image from the dockerfile path instead of using the predefined - image. Either docker_path or image must be set. - input_dir: The directory in the container where input files will be placed. - output_dir: The directory in the container where output files will be - retrieved from. - """ + """A code executor that uses a custom container to execute code. + + Attributes: + base_url: Optional. The base url of the user hosted Docker client. + image: The tag of the predefined image or custom image to run on the + container. Either docker_path or image must be set. + docker_path: The path to the directory containing the Dockerfile. If set, + build the image from the dockerfile path instead of using the predefined + image. Either docker_path or image must be set. + input_dir: The directory in the container where input files will be placed. + output_dir: The directory in the container where output files will be + retrieved from. + """ - base_url: Optional[str] = None - """ + base_url: Optional[str] = None + """ Optional. The base url of the user hosted Docker client. """ - image: str = None - """ + image: str = None + """ The tag of the predefined image or custom image to run on the container. Either docker_path or image must be set. """ - docker_path: str = None - """ + docker_path: str = None + """ The path to the directory containing the Dockerfile. If set, build the image from the dockerfile path instead of using the predefined image. Either docker_path or image must be set . """ - input_dir: str = "/tmp/inputs" - """ + input_dir: str = "/tmp/inputs" + """ The directory in the container where input files will be placed. """ - output_dir: str = "/tmp/outputs" - """ + output_dir: str = "/tmp/outputs" + """ The directory in the container where output files will be retrieved from. """ - stateful: bool = Field(default=False, frozen=True, exclude=True) + stateful: bool = Field(default=False, frozen=True, exclude=True) + + optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) - optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) + _client: DockerClient = None + _container: Container = None - _client: DockerClient = None - _container: Container = None + def __init__( + self, + base_url: Optional[str] = None, + image: Optional[str] = None, + docker_path: Optional[str] = None, + input_dir: Optional[str] = None, + output_dir: Optional[str] = None, + **data, + ): + """Initializes the ContainerCodeExecutor. - def __init__( - self, - base_url: Optional[str] = None, - image: Optional[str] = None, - docker_path: Optional[str] = None, - input_dir: Optional[str] = None, - output_dir: Optional[str] = None, - **data, + Args: + base_url: Optional. The base url of the user hosted Docker client. + image: The tag of the predefined image or custom image to run on the + container. Either docker_path or image must be set. + docker_path: The path to the directory containing the Dockerfile. If set, + build the image from the dockerfile path instead of using the predefined + image. Either docker_path or image must be set. + input_dir: The directory in the container where input files will be placed. + Defaults to '/tmp/inputs'. + output_dir: The directory in the container where output files will be + retrieved from. Defaults to '/tmp/outputs'. + **data: The data to initialize the ContainerCodeExecutor. + """ + if not image and not docker_path: + raise ValueError( + "Either image or docker_path must be set for ContainerCodeExecutor." + ) + if "stateful" in data and data["stateful"]: + raise ValueError("Cannot set `stateful=True` in ContainerCodeExecutor.") + if "optimize_data_file" in data and data["optimize_data_file"]: + raise ValueError( + "Cannot set `optimize_data_file=True` in ContainerCodeExecutor." + ) + + super().__init__(**data) + self.base_url = base_url + self.image = image if image else DEFAULT_IMAGE_TAG + self.docker_path = os.path.abspath(docker_path) if docker_path else None + self.input_dir = input_dir if input_dir else "/tmp/inputs" + self.output_dir = output_dir if output_dir else "/tmp/outputs" + + self._client = ( + docker.from_env() + if not self.base_url + else docker.DockerClient(base_url=self.base_url) + ) + self.__init_container() + + atexit.register(self.__cleanup_container) + + @override + def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + if code_execution_input.input_files: + self._put_input_files(code_execution_input.input_files) + + self._create_output_directory() + + output = "" + error = "" + exec_result = self._container.exec_run( + ["python3", "-c", code_execution_input.code], + demux=True, + ) + logger.debug("Executed code:\n```\n%s\n```", code_execution_input.code) + + if exec_result.output and exec_result.output[0]: + output = exec_result.output[0].decode("utf-8") + if ( + exec_result.output + and len(exec_result.output) > 1 + and exec_result.output[1] ): - """Initializes the ContainerCodeExecutor. - - Args: - base_url: Optional. The base url of the user hosted Docker client. - image: The tag of the predefined image or custom image to run on the - container. Either docker_path or image must be set. - docker_path: The path to the directory containing the Dockerfile. If set, - build the image from the dockerfile path instead of using the predefined - image. Either docker_path or image must be set. - input_dir: The directory in the container where input files will be placed. - Defaults to '/tmp/inputs'. - output_dir: The directory in the container where output files will be - retrieved from. Defaults to '/tmp/outputs'. - **data: The data to initialize the ContainerCodeExecutor. - """ - if not image and not docker_path: - raise ValueError( - "Either image or docker_path must be set for ContainerCodeExecutor." - ) - if "stateful" in data and data["stateful"]: - raise ValueError("Cannot set `stateful=True` in ContainerCodeExecutor.") - if "optimize_data_file" in data and data["optimize_data_file"]: - raise ValueError( - "Cannot set `optimize_data_file=True` in ContainerCodeExecutor." - ) - - super().__init__(**data) - self.base_url = base_url - self.image = image if image else DEFAULT_IMAGE_TAG - self.docker_path = os.path.abspath(docker_path) if docker_path else None - self.input_dir = input_dir if input_dir else "/tmp/inputs" - self.output_dir = output_dir if output_dir else "/tmp/outputs" - - self._client = ( - docker.from_env() - if not self.base_url - else docker.DockerClient(base_url=self.base_url) - ) - self.__init_container() - - atexit.register(self.__cleanup_container) - - @override - def execute_code( - self, - invocation_context: InvocationContext, - code_execution_input: CodeExecutionInput, - ) -> CodeExecutionResult: - if code_execution_input.input_files: - self._put_input_files(code_execution_input.input_files) - - self._create_output_directory() - - output = "" - error = "" - exec_result = self._container.exec_run( - ["python3", "-c", code_execution_input.code], - demux=True, - ) - logger.debug("Executed code:\n```\n%s\n```", code_execution_input.code) - - if exec_result.output and exec_result.output[0]: - output = exec_result.output[0].decode("utf-8") - if exec_result.output and len(exec_result.output) > 1 and exec_result.output[1]: - error = exec_result.output[1].decode("utf-8") - - output_files = self._get_output_files() - - return CodeExecutionResult( - stdout=output, - stderr=error, - output_files=output_files, - ) - - def _build_docker_image(self): - """Builds the Docker image.""" - if not self.docker_path: - raise ValueError("Docker path is not set.") - if not os.path.exists(self.docker_path): - raise FileNotFoundError(f"Invalid Docker path: {self.docker_path}") - - logger.info("Building Docker image...") - self._client.images.build( - path=self.docker_path, - tag=self.image, - rm=True, - ) - logger.info("Docker image: %s built.", self.image) - - def _verify_python_installation(self): - """Verifies the container has python3 installed.""" - exec_result = self._container.exec_run(["which", "python3"]) - if exec_result.exit_code != 0: - raise ValueError("python3 is not installed in the container.") - - def _put_input_files(self, input_files: list[File]) -> None: - """Puts input files into the container using put_archive. - - Args: - input_files: The list of input files to copy into the container. - """ - tar_buffer = io.BytesIO() - with tarfile.open(fileobj=tar_buffer, mode="w") as tar: - for file in input_files: - content = file.content - if isinstance(content, str): - content = content.encode("utf-8") - tarinfo = tarfile.TarInfo(name=file.name) - tarinfo.size = len(content) - tar.addfile(tarinfo, io.BytesIO(content)) - - tar_buffer.seek(0) - self._container.put_archive( - self.input_dir, - tar_buffer.read(), - ) - logger.debug("Copied %d input files to %s", len(input_files), self.input_dir) - - def _create_output_directory(self) -> None: - """Creates the output directory in the container if it doesn't exist.""" - exec_result = self._container.exec_run( - ["mkdir", "-p", self.output_dir], - ) - if exec_result.exit_code != 0: - logger.warning( - "Failed to create output directory %s: %s", - self.output_dir, - exec_result.output, - ) - - def _get_output_files(self) -> list[File]: - """Gets output files from the container. - - Returns: - The list of output files retrieved from the container. - """ - try: - tar_bytes, stat = self._container.get_archive(self.output_dir) - except docker.errors.APIError as e: - if e.response.status_code == 404: - logger.debug("No output files found at %s", self.output_dir) - return [] - raise - - tar_buffer = io.BytesIO(tar_bytes) - output_files = [] - - with tarfile.open(fileobj=tar_buffer, mode="r") as tar: - for member in tar.getmembers(): - if member.isfile(): - file_obj = tar.extractfile(member) - if file_obj: - content = file_obj.read() - file_name = os.path.basename(member.name) - if file_name: - output_files.append( - File( - name=file_name, - content=content, - mime_type=self._guess_mime_type(file_name), - ) - ) - - logger.debug( - "Retrieved %d output files from %s", len(output_files), self.output_dir - ) - return output_files - - def _guess_mime_type(self, file_name: str) -> str: - """Guesses the MIME type based on the file extension. - - Args: - file_name: The name of the file. - - Returns: - The guessed MIME type, or 'application/octet-stream' if unknown. - """ - import mimetypes - - mime_type, _ = mimetypes.guess_type(file_name) - return mime_type if mime_type else "application/octet-stream" - - def __init_container(self): - """Initializes the container.""" - if not self._client: - raise RuntimeError("Docker client is not initialized.") - - if self.docker_path: - self._build_docker_image() - - logger.info("Starting container for ContainerCodeExecutor...") - self._container = self._client.containers.run( - image=self.image, - detach=True, - tty=True, - ) - logger.info("Container %s started.", self._container.id) - - self._verify_python_installation() - - def __cleanup_container(self): - """Closes the container on exit.""" - if not self._container: - return - - logger.info("[Cleanup] Stopping the container...") - self._container.stop() - self._container.remove() - logger.info("Container %s stopped and removed.", self._container.id) + error = exec_result.output[1].decode("utf-8") + + output_files = self._get_output_files() + + return CodeExecutionResult( + stdout=output, + stderr=error, + output_files=output_files, + ) + + def _build_docker_image(self): + """Builds the Docker image.""" + if not self.docker_path: + raise ValueError("Docker path is not set.") + if not os.path.exists(self.docker_path): + raise FileNotFoundError(f"Invalid Docker path: {self.docker_path}") + + logger.info("Building Docker image...") + self._client.images.build( + path=self.docker_path, + tag=self.image, + rm=True, + ) + logger.info("Docker image: %s built.", self.image) + + def _verify_python_installation(self): + """Verifies the container has python3 installed.""" + exec_result = self._container.exec_run(["which", "python3"]) + if exec_result.exit_code != 0: + raise ValueError("python3 is not installed in the container.") + + def _put_input_files(self, input_files: list[File]) -> None: + """Puts input files into the container using put_archive. + + Args: + input_files: The list of input files to copy into the container. + """ + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + for file in input_files: + content = file.content + if isinstance(content, str): + content = content.encode("utf-8") + tarinfo = tarfile.TarInfo(name=file.name) + tarinfo.size = len(content) + tar.addfile(tarinfo, io.BytesIO(content)) + + tar_buffer.seek(0) + self._container.put_archive( + self.input_dir, + tar_buffer.read(), + ) + logger.debug( + "Copied %d input files to %s", len(input_files), self.input_dir + ) + + def _create_output_directory(self) -> None: + """Creates the output directory in the container if it doesn't exist.""" + exec_result = self._container.exec_run( + ["mkdir", "-p", self.output_dir], + ) + if exec_result.exit_code != 0: + logger.warning( + "Failed to create output directory %s: %s", + self.output_dir, + exec_result.output, + ) + + def _get_output_files(self) -> list[File]: + """Gets output files from the container. + + Returns: + The list of output files retrieved from the container. + """ + try: + tar_bytes, stat = self._container.get_archive(self.output_dir) + except docker.errors.APIError as e: + if e.response.status_code == 404: + logger.debug("No output files found at %s", self.output_dir) + return [] + raise + + tar_buffer = io.BytesIO(tar_bytes) + output_files = [] + + with tarfile.open(fileobj=tar_buffer, mode="r") as tar: + for member in tar.getmembers(): + if member.isfile(): + file_obj = tar.extractfile(member) + if file_obj: + content = file_obj.read() + file_name = os.path.basename(member.name) + if file_name: + output_files.append( + File( + name=file_name, + content=content, + mime_type=self._guess_mime_type(file_name), + ) + ) + + logger.debug( + "Retrieved %d output files from %s", len(output_files), self.output_dir + ) + return output_files + + def _guess_mime_type(self, file_name: str) -> str: + """Guesses the MIME type based on the file extension. + + Args: + file_name: The name of the file. + + Returns: + The guessed MIME type, or 'application/octet-stream' if unknown. + """ + import mimetypes + + mime_type, _ = mimetypes.guess_type(file_name) + return mime_type if mime_type else "application/octet-stream" + + def __init_container(self): + """Initializes the container.""" + if not self._client: + raise RuntimeError("Docker client is not initialized.") + + if self.docker_path: + self._build_docker_image() + + logger.info("Starting container for ContainerCodeExecutor...") + self._container = self._client.containers.run( + image=self.image, + detach=True, + tty=True, + ) + logger.info("Container %s started.", self._container.id) + + self._verify_python_installation() + + def __cleanup_container(self): + """Closes the container on exit.""" + if not self._container: + return + + logger.info("[Cleanup] Stopping the container...") + self._container.stop() + self._container.remove() + logger.info("Container %s stopped and removed.", self._container.id) From c6d1a1caa1fff6f47c92af31f579c97a90e10af6 Mon Sep 17 00:00:00 2001 From: Hamdi Sakly Date: Fri, 17 Apr 2026 21:23:00 +0100 Subject: [PATCH 5/7] add docker to test in pyproject.toml --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5815b971f5..66db38897a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,8 +123,8 @@ test = [ "a2a-sdk>=0.3.0,<0.4.0", "anthropic>=0.43.0", # For anthropic model tests "crewai[tools];python_version>='3.11' and python_version<'3.12'", # For CrewaiTool tests; chromadb/pypika fail on 3.12+ - "google-cloud-firestore>=2.11.0, <3.0.0", - "google-cloud-iamconnectorcredentials>=0.1.0, <0.2.0", + "docker>=7.0.0", + "google-cloud-firestore>=2.11.0", "google-cloud-parametermanager>=0.4.0, <1.0.0", "kubernetes>=29.0.0", # For GkeCodeExecutor "langchain-community>=0.3.17", From 49a1e8d292cc212c86ee63ea8ab416c6e76ab17d Mon Sep 17 00:00:00 2001 From: Hamdi Sakly Date: Sat, 18 Apr 2026 02:17:43 +0100 Subject: [PATCH 6/7] fix the testing of container execution and fix pyproject.toml file --- pyproject.toml | 1 + .../test_container_code_executor.py | 385 +++++++++--------- 2 files changed, 200 insertions(+), 186 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 66db38897a..9a37dc492c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,6 +123,7 @@ test = [ "a2a-sdk>=0.3.0,<0.4.0", "anthropic>=0.43.0", # For anthropic model tests "crewai[tools];python_version>='3.11' and python_version<'3.12'", # For CrewaiTool tests; chromadb/pypika fail on 3.12+ + "google-cloud-iamconnectorcredentials>=0.1.0, <0.2.0", "docker>=7.0.0", "google-cloud-firestore>=2.11.0", "google-cloud-parametermanager>=0.4.0, <1.0.0", diff --git a/tests/unittests/code_executors/test_container_code_executor.py b/tests/unittests/code_executors/test_container_code_executor.py index b49c186134..a96312ace6 100644 --- a/tests/unittests/code_executors/test_container_code_executor.py +++ b/tests/unittests/code_executors/test_container_code_executor.py @@ -28,230 +28,243 @@ @pytest.fixture def mock_container(): - container = mock.MagicMock() - container.id = "test-container-id" - container.exec_run.return_value = mock.MagicMock( - output=(b"Hello World", b""), - exit_code=0, - ) - return container + container = mock.MagicMock() + container.id = "test-container-id" + container.exec_run.return_value = mock.MagicMock( + output=(b"Hello World", b""), + exit_code=0, + ) + container.get_archive.return_value = (b"", mock.MagicMock()) + return container @pytest.fixture def mock_docker_client(mock_container): - client = mock.MagicMock() - client.containers.run.return_value = mock_container - return client + client = mock.MagicMock() + client.containers.run.return_value = mock_container + return client class TestContainerCodeExecutorInit: - - def test_init_with_image(self, mock_docker_client): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - assert executor.image == "test-image:latest" - assert executor.input_dir == "/tmp/inputs" - assert executor.output_dir == "/tmp/outputs" - - def test_init_with_custom_dirs(self, mock_docker_client): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor( - image="test-image:latest", - input_dir="/custom/inputs", - output_dir="/custom/outputs", - ) - assert executor.input_dir == "/custom/inputs" - assert executor.output_dir == "/custom/outputs" - - def test_init_requires_image_or_docker_path(self): - with pytest.raises( - ValueError, match="Either image or docker_path must be set" - ): - ContainerCodeExecutor() - - def test_init_rejects_stateful(self): - with pytest.raises(ValueError, match="Cannot set `stateful=True`"): - ContainerCodeExecutor(image="test", stateful=True) - - def test_init_rejects_optimize_data_file(self): - with pytest.raises( - ValueError, match="Cannot set `optimize_data_file=True`" - ): - ContainerCodeExecutor(image="test", optimize_data_file=True) + def test_init_with_image(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + assert executor.image == "test-image:latest" + assert executor.input_dir == "/tmp/inputs" + assert executor.output_dir == "/tmp/outputs" + + def test_init_with_custom_dirs(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor( + image="test-image:latest", + input_dir="/custom/inputs", + output_dir="/custom/outputs", + ) + assert executor.input_dir == "/custom/inputs" + assert executor.output_dir == "/custom/outputs" + + def test_init_requires_image_or_docker_path(self): + with pytest.raises(ValueError, match="Either image or docker_path must be set"): + ContainerCodeExecutor() + + def test_init_rejects_stateful(self): + with pytest.raises(ValueError, match="Cannot set `stateful=True`"): + ContainerCodeExecutor(image="test", stateful=True) + + def test_init_rejects_optimize_data_file(self): + with pytest.raises(ValueError, match="Cannot set `optimize_data_file=True`"): + ContainerCodeExecutor(image="test", optimize_data_file=True) class TestExecuteCode: - - def test_execute_code_basic(self, mock_docker_client, mock_container): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput(code='print("Hello World")') - - result = executor.execute_code(context, code_input) - - assert result.stdout == "Hello World" - assert result.stderr == "" - assert result.output_files == [] - - def test_execute_code_with_error(self, mock_docker_client, mock_container): - mock_container.exec_run.return_value = mock.MagicMock( - output=(b"", b"Some error"), - exit_code=1, - ) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput(code='raise Error("test")') - - result = executor.execute_code(context, code_input) - - assert result.stderr == "Some error" - - def test_execute_code_with_input_files( - self, mock_docker_client, mock_container - ): - mock_container.put_archive = mock.MagicMock() - mock_container.exec_run.return_value = mock.MagicMock( - output=(b"", b""), - exit_code=0, - ) - mock_container.get_archive = mock.MagicMock( - side_effect=mock.MagicMock(response=mock.MagicMock(status_code=404)), - ) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput( - code='print("test")', - input_files=[File(name="test.txt", content="test content")], - ) - - result = executor.execute_code(context, code_input) - - mock_container.put_archive.assert_called_once() - call_args = mock_container.put_archive.call_args - assert call_args[0][0] == "/tmp/inputs" - - def test_execute_code_with_output_files( - self, mock_docker_client, mock_container - ): - mock_container.exec_run.return_value = mock.MagicMock( - output=(b"", b""), - exit_code=0, - ) - - tar_buffer = io.BytesIO() - with tarfile.open(fileobj=tar_buffer, mode="w") as tar: - tar.addfile( - tarfile.TarInfo(name="output.txt"), - io.BytesIO(b"output content"), - ) - tar_buffer.seek(0) - tar_bytes = tar_buffer.read() - - mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) - - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput(code='print("test")') - - result = executor.execute_code(context, code_input) - - assert len(result.output_files) == 1 - assert result.output_files[0].name == "output.txt" + def test_execute_code_basic(self, mock_docker_client, mock_container): + import docker + + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("Hello World")') + + result = executor.execute_code(context, code_input) + + assert result.stdout == "Hello World" + assert result.stderr == "" + assert result.output_files == [] + + def test_execute_code_with_error(self, mock_docker_client, mock_container): + import docker + + call_count = [0] + + def exec_run_side_effect(cmd, demux=False): + call_count[0] += 1 + if call_count[0] == 3: + return mock.MagicMock( + exit_code=1, + output=(b"", b"Some error"), + ) + return mock.MagicMock(exit_code=0, output=(b"", b"")) + + mock_container.exec_run.side_effect = exec_run_side_effect + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='raise Error("test")') + + result = executor.execute_code(context, code_input) + + assert result.stderr == "Some error" + + def test_execute_code_with_input_files(self, mock_docker_client, mock_container): + import docker + + mock_container.put_archive = mock.MagicMock() + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput( + code='print("test")', + input_files=[File(name="test.txt", content="test content")], + ) + + result = executor.execute_code(context, code_input) + + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" + + def test_execute_code_with_output_files(self, mock_docker_client, mock_container): + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + + content = b"output content" + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar_info = tarfile.TarInfo(name="output.txt") + tar_info.size = len(content) + tar.addfile(tar_info, io.BytesIO(content)) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() + + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("test")') + + result = executor.execute_code(context, code_input) + + assert len(result.output_files) == 1 + assert result.output_files[0].name == "output.txt" class TestPutInputFiles: + def test_put_archive_called(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [ + File(name="file1.txt", content="content1"), + File(name="file2.txt", content="content2"), + ] - def test_put_archive_called(self, mock_docker_client, mock_container): - mock_container.put_archive = mock.MagicMock() - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - input_files = [ - File(name="file1.txt", content="content1"), - File(name="file2.txt", content="content2"), - ] + executor._put_input_files(input_files) - executor._put_input_files(input_files) + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" - mock_container.put_archive.assert_called_once() - call_args = mock_container.put_archive.call_args - assert call_args[0][0] == "/tmp/inputs" + def test_handles_string_content(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [File(name="test.txt", content="string content")] - def test_handles_string_content(self, mock_docker_client, mock_container): - mock_container.put_archive = mock.MagicMock() - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - input_files = [File(name="test.txt", content="string content")] + executor._put_input_files(input_files) - executor._put_input_files(input_files) - - mock_container.put_archive.assert_called_once() + mock_container.put_archive.assert_called_once() class TestGetOutputFiles: + def test_no_output_files(self, mock_docker_client, mock_container): + import docker - def test_no_output_files(self, mock_docker_client, mock_container): - import docker - - mock_container.get_archive.side_effect = docker.errors.APIError( - "Not found", response=mock.MagicMock(status_code=404) - ) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - output_files = executor._get_output_files() + output_files = executor._get_output_files() - assert output_files == [] + assert output_files == [] - def test_extracts_files_from_archive( - self, mock_docker_client, mock_container - ): - tar_buffer = io.BytesIO() - with tarfile.open(fileobj=tar_buffer, mode="w") as tar: - tar.addfile( - tarfile.TarInfo(name="output.txt"), - io.BytesIO(b"output content"), - ) - tar_buffer.seek(0) - tar_bytes = tar_buffer.read() + def test_extracts_files_from_archive(self, mock_docker_client, mock_container): + content = b"output content" + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar_info = tarfile.TarInfo(name="output.txt") + tar_info.size = len(content) + tar.addfile(tar_info, io.BytesIO(content)) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() - mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - output_files = executor._get_output_files() + output_files = executor._get_output_files() - assert len(output_files) == 1 - assert output_files[0].name == "output.txt" - assert output_files[0].content == b"output content" + assert len(output_files) == 1 + assert output_files[0].name == "output.txt" + assert output_files[0].content == content class TestMimeTypeGuessing: + def test_guess_txt(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - def test_guess_txt(self, mock_docker_client, mock_container): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + mime_type = executor._guess_mime_type("test.txt") - mime_type = executor._guess_mime_type("test.txt") + assert mime_type == "text/plain" - assert mime_type == "text/plain" + def test_guess_csv(self, mock_docker_client, mock_container): + import mimetypes as mimetypes_module - def test_guess_csv(self, mock_docker_client, mock_container): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + original_guess_type = mimetypes_module.guess_type + mimetypes_module.guess_type = lambda f: ("text/csv", None) + try: + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - mime_type = executor._guess_mime_type("data.csv") + mime_type = executor._guess_mime_type("data.csv") - assert mime_type == "text/csv" + assert mime_type == "text/csv" + finally: + mimetypes_module.guess_type = original_guess_type - def test_default_for_unknown(self, mock_docker_client, mock_container): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + def test_default_for_unknown(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - mime_type = executor._guess_mime_type("unknown.xyz") + mime_type = executor._guess_mime_type("unknown.xyz") - assert mime_type == "application/octet-stream" + assert mime_type == "application/octet-stream" From 7abef2601489ac59bd919a0d7417a08fa7462066 Mon Sep 17 00:00:00 2001 From: Hamdi Sakly Date: Sun, 19 Apr 2026 14:21:00 +0100 Subject: [PATCH 7/7] chaged the test file and changed the optimize_data_file to true --- .../code_executors/container_code_executor.py | 13 +- .../test_container_code_executor.py | 392 +++++++++--------- 2 files changed, 210 insertions(+), 195 deletions(-) diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index d301b32414..ccfeebe253 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -82,7 +82,7 @@ class ContainerCodeExecutor(BaseCodeExecutor): stateful: bool = Field(default=False, frozen=True, exclude=True) - optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) + optimize_data_file: bool = Field(default=True, frozen=True, exclude=True) _client: DockerClient = None _container: Container = None @@ -117,10 +117,6 @@ def __init__( ) if "stateful" in data and data["stateful"]: raise ValueError("Cannot set `stateful=True` in ContainerCodeExecutor.") - if "optimize_data_file" in data and data["optimize_data_file"]: - raise ValueError( - "Cannot set `optimize_data_file=True` in ContainerCodeExecutor." - ) super().__init__(**data) self.base_url = base_url @@ -239,13 +235,16 @@ def _get_output_files(self) -> list[File]: The list of output files retrieved from the container. """ try: - tar_bytes, stat = self._container.get_archive(self.output_dir) + tar_stream, _ = self._container.get_archive(self.output_dir) except docker.errors.APIError as e: if e.response.status_code == 404: logger.debug("No output files found at %s", self.output_dir) return [] raise - + if isinstance(tar_stream, bytes): + tar_bytes = tar_stream + else: + tar_bytes = b"".join(tar_stream) tar_buffer = io.BytesIO(tar_bytes) output_files = [] diff --git a/tests/unittests/code_executors/test_container_code_executor.py b/tests/unittests/code_executors/test_container_code_executor.py index a96312ace6..fabfaf62e1 100644 --- a/tests/unittests/code_executors/test_container_code_executor.py +++ b/tests/unittests/code_executors/test_container_code_executor.py @@ -28,243 +28,259 @@ @pytest.fixture def mock_container(): - container = mock.MagicMock() - container.id = "test-container-id" - container.exec_run.return_value = mock.MagicMock( - output=(b"Hello World", b""), - exit_code=0, - ) - container.get_archive.return_value = (b"", mock.MagicMock()) - return container + container = mock.MagicMock() + container.id = "test-container-id" + container.exec_run.return_value = mock.MagicMock( + output=(b"Hello World", b""), + exit_code=0, + ) + container.get_archive.return_value = (b"", mock.MagicMock()) + return container @pytest.fixture def mock_docker_client(mock_container): - client = mock.MagicMock() - client.containers.run.return_value = mock_container - return client + client = mock.MagicMock() + client.containers.run.return_value = mock_container + return client class TestContainerCodeExecutorInit: - def test_init_with_image(self, mock_docker_client): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - assert executor.image == "test-image:latest" - assert executor.input_dir == "/tmp/inputs" - assert executor.output_dir == "/tmp/outputs" - - def test_init_with_custom_dirs(self, mock_docker_client): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor( - image="test-image:latest", - input_dir="/custom/inputs", - output_dir="/custom/outputs", - ) - assert executor.input_dir == "/custom/inputs" - assert executor.output_dir == "/custom/outputs" - - def test_init_requires_image_or_docker_path(self): - with pytest.raises(ValueError, match="Either image or docker_path must be set"): - ContainerCodeExecutor() - - def test_init_rejects_stateful(self): - with pytest.raises(ValueError, match="Cannot set `stateful=True`"): - ContainerCodeExecutor(image="test", stateful=True) - - def test_init_rejects_optimize_data_file(self): - with pytest.raises(ValueError, match="Cannot set `optimize_data_file=True`"): - ContainerCodeExecutor(image="test", optimize_data_file=True) + + def test_init_with_image(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + assert executor.image == "test-image:latest" + assert executor.input_dir == "/tmp/inputs" + assert executor.output_dir == "/tmp/outputs" + + def test_init_with_custom_dirs(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor( + image="test-image:latest", + input_dir="/custom/inputs", + output_dir="/custom/outputs", + ) + assert executor.input_dir == "/custom/inputs" + assert executor.output_dir == "/custom/outputs" + + def test_init_requires_image_or_docker_path(self): + with pytest.raises( + ValueError, match="Either image or docker_path must be set" + ): + ContainerCodeExecutor() + + def test_init_rejects_stateful(self): + with pytest.raises(ValueError, match="Cannot set `stateful=True`"): + ContainerCodeExecutor(image="test", stateful=True) + + def test_init_allows_optimize_data_file(self, mock_docker_client): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor( + image="test-image:latest", optimize_data_file=True + ) + assert executor.optimize_data_file is True class TestExecuteCode: - def test_execute_code_basic(self, mock_docker_client, mock_container): - import docker - mock_container.get_archive.side_effect = docker.errors.APIError( - "Not found", response=mock.MagicMock(status_code=404) - ) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput(code='print("Hello World")') - - result = executor.execute_code(context, code_input) - - assert result.stdout == "Hello World" - assert result.stderr == "" - assert result.output_files == [] - - def test_execute_code_with_error(self, mock_docker_client, mock_container): - import docker - - call_count = [0] - - def exec_run_side_effect(cmd, demux=False): - call_count[0] += 1 - if call_count[0] == 3: - return mock.MagicMock( - exit_code=1, - output=(b"", b"Some error"), - ) - return mock.MagicMock(exit_code=0, output=(b"", b"")) - - mock_container.exec_run.side_effect = exec_run_side_effect - mock_container.get_archive.side_effect = docker.errors.APIError( - "Not found", response=mock.MagicMock(status_code=404) - ) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput(code='raise Error("test")') + def test_execute_code_basic(self, mock_docker_client, mock_container): + import docker - result = executor.execute_code(context, code_input) + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("Hello World")') - assert result.stderr == "Some error" + result = executor.execute_code(context, code_input) - def test_execute_code_with_input_files(self, mock_docker_client, mock_container): - import docker + assert result.stdout == "Hello World" + assert result.stderr == "" + assert result.output_files == [] - mock_container.put_archive = mock.MagicMock() - mock_container.exec_run.return_value = mock.MagicMock( - output=(b"", b""), - exit_code=0, - ) - mock_container.get_archive.side_effect = docker.errors.APIError( - "Not found", response=mock.MagicMock(status_code=404) - ) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput( - code='print("test")', - input_files=[File(name="test.txt", content="test content")], - ) - - result = executor.execute_code(context, code_input) - - mock_container.put_archive.assert_called_once() - call_args = mock_container.put_archive.call_args - assert call_args[0][0] == "/tmp/inputs" - - def test_execute_code_with_output_files(self, mock_docker_client, mock_container): - mock_container.exec_run.return_value = mock.MagicMock( - output=(b"", b""), - exit_code=0, + def test_execute_code_with_error(self, mock_docker_client, mock_container): + import docker + + call_count = [0] + + def exec_run_side_effect(cmd, demux=False): + call_count[0] += 1 + if call_count[0] == 3: + return mock.MagicMock( + exit_code=1, + output=(b"", b"Some error"), ) + return mock.MagicMock(exit_code=0, output=(b"", b"")) + + mock_container.exec_run.side_effect = exec_run_side_effect + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='raise Error("test")') + + result = executor.execute_code(context, code_input) - content = b"output content" - tar_buffer = io.BytesIO() - with tarfile.open(fileobj=tar_buffer, mode="w") as tar: - tar_info = tarfile.TarInfo(name="output.txt") - tar_info.size = len(content) - tar.addfile(tar_info, io.BytesIO(content)) - tar_buffer.seek(0) - tar_bytes = tar_buffer.read() + assert result.stderr == "Some error" - mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + def test_execute_code_with_input_files( + self, mock_docker_client, mock_container + ): + import docker - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - context = mock.MagicMock(spec=InvocationContext) - code_input = CodeExecutionInput(code='print("test")') + mock_container.put_archive = mock.MagicMock() + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput( + code='print("test")', + input_files=[File(name="test.txt", content="test content")], + ) + + result = executor.execute_code(context, code_input) + + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" + + def test_execute_code_with_output_files( + self, mock_docker_client, mock_container + ): + mock_container.exec_run.return_value = mock.MagicMock( + output=(b"", b""), + exit_code=0, + ) + + content = b"output content" + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar_info = tarfile.TarInfo(name="output.txt") + tar_info.size = len(content) + tar.addfile(tar_info, io.BytesIO(content)) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() + + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) - result = executor.execute_code(context, code_input) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + context = mock.MagicMock(spec=InvocationContext) + code_input = CodeExecutionInput(code='print("test")') - assert len(result.output_files) == 1 - assert result.output_files[0].name == "output.txt" + result = executor.execute_code(context, code_input) + + assert len(result.output_files) == 1 + assert result.output_files[0].name == "output.txt" class TestPutInputFiles: - def test_put_archive_called(self, mock_docker_client, mock_container): - mock_container.put_archive = mock.MagicMock() - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - input_files = [ - File(name="file1.txt", content="content1"), - File(name="file2.txt", content="content2"), - ] - executor._put_input_files(input_files) + def test_put_archive_called(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [ + File(name="file1.txt", content="content1"), + File(name="file2.txt", content="content2"), + ] + + executor._put_input_files(input_files) - mock_container.put_archive.assert_called_once() - call_args = mock_container.put_archive.call_args - assert call_args[0][0] == "/tmp/inputs" + mock_container.put_archive.assert_called_once() + call_args = mock_container.put_archive.call_args + assert call_args[0][0] == "/tmp/inputs" - def test_handles_string_content(self, mock_docker_client, mock_container): - mock_container.put_archive = mock.MagicMock() - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - input_files = [File(name="test.txt", content="string content")] + def test_handles_string_content(self, mock_docker_client, mock_container): + mock_container.put_archive = mock.MagicMock() + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + input_files = [File(name="test.txt", content="string content")] - executor._put_input_files(input_files) + executor._put_input_files(input_files) - mock_container.put_archive.assert_called_once() + mock_container.put_archive.assert_called_once() class TestGetOutputFiles: - def test_no_output_files(self, mock_docker_client, mock_container): - import docker - mock_container.get_archive.side_effect = docker.errors.APIError( - "Not found", response=mock.MagicMock(status_code=404) - ) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + def test_no_output_files(self, mock_docker_client, mock_container): + import docker - output_files = executor._get_output_files() + mock_container.get_archive.side_effect = docker.errors.APIError( + "Not found", response=mock.MagicMock(status_code=404) + ) + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + output_files = executor._get_output_files() - assert output_files == [] + assert output_files == [] - def test_extracts_files_from_archive(self, mock_docker_client, mock_container): - content = b"output content" - tar_buffer = io.BytesIO() - with tarfile.open(fileobj=tar_buffer, mode="w") as tar: - tar_info = tarfile.TarInfo(name="output.txt") - tar_info.size = len(content) - tar.addfile(tar_info, io.BytesIO(content)) - tar_buffer.seek(0) - tar_bytes = tar_buffer.read() + def test_extracts_files_from_archive( + self, mock_docker_client, mock_container + ): + content = b"output content" + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + tar_info = tarfile.TarInfo(name="output.txt") + tar_info.size = len(content) + tar.addfile(tar_info, io.BytesIO(content)) + tar_buffer.seek(0) + tar_bytes = tar_buffer.read() - mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) + mock_container.get_archive.return_value = (tar_bytes, mock.MagicMock()) - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - output_files = executor._get_output_files() + output_files = executor._get_output_files() - assert len(output_files) == 1 - assert output_files[0].name == "output.txt" - assert output_files[0].content == content + assert len(output_files) == 1 + assert output_files[0].name == "output.txt" + assert output_files[0].content == content class TestMimeTypeGuessing: - def test_guess_txt(self, mock_docker_client, mock_container): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") - mime_type = executor._guess_mime_type("test.txt") + def test_guess_txt(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") + + mime_type = executor._guess_mime_type("test.txt") - assert mime_type == "text/plain" + assert mime_type == "text/plain" - def test_guess_csv(self, mock_docker_client, mock_container): - import mimetypes as mimetypes_module + def test_guess_csv(self, mock_docker_client, mock_container): + import mimetypes as mimetypes_module - original_guess_type = mimetypes_module.guess_type - mimetypes_module.guess_type = lambda f: ("text/csv", None) - try: - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + original_guess_type = mimetypes_module.guess_type + mimetypes_module.guess_type = lambda f: ("text/csv", None) + try: + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - mime_type = executor._guess_mime_type("data.csv") + mime_type = executor._guess_mime_type("data.csv") - assert mime_type == "text/csv" - finally: - mimetypes_module.guess_type = original_guess_type + assert mime_type == "text/csv" + finally: + mimetypes_module.guess_type = original_guess_type - def test_default_for_unknown(self, mock_docker_client, mock_container): - with mock.patch("docker.from_env", return_value=mock_docker_client): - executor = ContainerCodeExecutor(image="test-image:latest") + def test_default_for_unknown(self, mock_docker_client, mock_container): + with mock.patch("docker.from_env", return_value=mock_docker_client): + executor = ContainerCodeExecutor(image="test-image:latest") - mime_type = executor._guess_mime_type("unknown.xyz") + mime_type = executor._guess_mime_type("unknown.xyz") - assert mime_type == "application/octet-stream" + assert mime_type == "application/octet-stream"