testcontainers · mgorsk1 · Aug 12, 2024 · Aug 12, 2024 · Aug 12, 2024 · Aug 12, 2024
diff --git 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
@@ -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
@@ -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