diff --git a/core/testcontainers/core/container.py b/core/testcontainers/core/container.py index e9415441..a90c726b 100644 --- a/core/testcontainers/core/container.py +++ b/core/testcontainers/core/container.py @@ -1,4 +1,7 @@ import contextlib +import io +import os +import tarfile from platform import system from socket import socket from typing import TYPE_CHECKING, Optional @@ -13,6 +16,7 @@ from testcontainers.core.exceptions import ContainerStartException from testcontainers.core.labels import LABEL_SESSION_ID, SESSION_ID from testcontainers.core.network import Network +from testcontainers.core.transferable import Transferable from testcontainers.core.utils import inside_container, is_arm, setup_logger from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs @@ -52,6 +56,7 @@ def __init__( self._network: Optional[Network] = None self._network_aliases: Optional[list[str]] = None self._kwargs = kwargs + self._files: list[Transferable] = [] def with_env(self, key: str, value: str) -> Self: self.env[key] = value @@ -78,6 +83,33 @@ def with_kwargs(self, **kwargs) -> Self: self._kwargs = kwargs return self + def with_copy_file_to_container(self, transferable: Transferable) -> Self: + self._files.append(transferable) + + return self + + def copy_file_from_container(self, container_file: os.PathLike, destination_file: os.PathLike) -> os.PathLike: + tar_stream, _ = self._container.get_archive(container_file) + + for chunk in tar_stream: + with tarfile.open(fileobj=io.BytesIO(chunk)) as tar: + for member in tar.getmembers(): + with open(destination_file, "wb") as f: + f.write(tar.extractfile(member).read()) + + return destination_file + + @staticmethod + def _put_data_in_container(container, transferable: Transferable): + data = io.BytesIO() + + with transferable as f, tarfile.open(fileobj=data, mode="w") as tar: + tar.add(f.input_path, arcname=f.output_path) + + data.seek(0) + + container.put_archive("/", data) + def maybe_emulate_amd64(self) -> Self: if is_arm(): return self.with_kwargs(platform="linux/amd64") @@ -115,6 +147,10 @@ def start(self) -> Self: ) logger.info("Container started: %s", self._container.short_id) + + for transferable in self._files: + DockerContainer._put_data_in_container(self._container, transferable) + return self def stop(self, force=True, delete_volume=True) -> None: diff --git a/core/testcontainers/core/transferable.py b/core/testcontainers/core/transferable.py new file mode 100644 index 00000000..6f6baab0 --- /dev/null +++ b/core/testcontainers/core/transferable.py @@ -0,0 +1,33 @@ +import os +import tempfile +from typing import Union + + +class Transferable: + def __init__(self, input_data: Union[os.PathLike, bytes], output_path: os.PathLike): + self._input = input_data + self._output_path = output_path + + self._tmp_file: bool = False + + def __enter__(self): + if isinstance(self._input, bytes): + tmp_file = tempfile.NamedTemporaryFile(delete=False) + tmp_file.write(self._input) + + self._input = tmp_file.name + self._tmp_file = True + + return self + + def __exit__(self, *args): + if self._tmp_file: + os.remove(self._input) + + @property + def input_path(self): + return self._input + + @property + def output_path(self): + return self._output_path diff --git a/core/tests/test_core.py b/core/tests/test_core.py index 8d0c7794..0329f975 100644 --- a/core/tests/test_core.py +++ b/core/tests/test_core.py @@ -8,6 +8,7 @@ from testcontainers.core.container import DockerContainer from testcontainers.core.image import DockerImage +from testcontainers.core.transferable import Transferable from testcontainers.core.waiting_utils import wait_for_logs @@ -92,3 +93,33 @@ def test_docker_image_with_custom_dockerfile_path(dockerfile_path: Optional[Path with DockerContainer(str(image)) as container: assert container._container.image.short_id.endswith(image_short_id), "Image ID mismatch" assert container.get_logs() == (("Hello world!\n").encode(), b""), "Container logs mismatch" + + +def test_docker_start_with_copy_file_to_container_from_binary_transferable() -> None: + container = DockerContainer("nginx") + data = "test_docker_start_with_copy_file_to_container_from_binary_transferable" + + input_data = data.encode("utf-8") + output_file = Path("/tmp/test_docker_start_with_copy_file_to_container_from_binary_transferable.txt") + + container.with_copy_file_to_container(Transferable(input_data, output_file)).start() + + _, stdout = container.exec(f"cat {output_file}") + assert stdout.decode() == data + + +def test_docker_start_with_copy_file_to_container_from_file_transferable() -> None: + container = DockerContainer("nginx") + data = "test_docker_start_with_copy_file_to_container_from_file_transferable" + + with tempfile.NamedTemporaryFile(delete=True) as f: + f.write(data.encode("utf-8")) + f.seek(0) + + input_file = Path(f.name) + output_file = Path("/tmp/test_docker_start_with_copy_file_to_container_from_file_transferable.txt") + + container.with_copy_file_to_container(Transferable(input_file, output_file)).start() + + _, stdout = container.exec(f"cat {output_file}") + assert stdout.decode() == data