Tests demonstrating copy file interfaces

Author: rhoban13

core/testcontainers/core/container.py

@@ -1,4 +1,8 @@
 import contextlib
+import dataclasses
+import io
+import pathlib
+import tarfile
 from os import PathLike
 from socket import socket
 from types import TracebackType
@@ -31,6 +35,13 @@ class Mount(TypedDict):
     mode: str
 
 
+@dataclasses.dataclass
+class Transferrable:
+    source: Union[bytes, PathLike]
+    destination_in_container: str
+    mode: int = 0o644
+
+
 class DockerContainer:
     """
     Basic container object to spin up Docker instances.
@@ -69,6 +80,7 @@ def __init__(
         volumes: Optional[list[tuple[str, str, str]]] = None,
         network: Optional[Network] = None,
         network_aliases: Optional[list[str]] = None,
+        transferrables: Optional[tuple[Transferrable]] = None,
         **kwargs: Any,
     ) -> None:
         self.env = env or {}
@@ -96,6 +108,7 @@ def __init__(
             self.with_network_aliases(*network_aliases)
 
         self._kwargs = kwargs
+        self._transferrables = transferrables or []
 
     def with_env(self, key: str, value: str) -> Self:
         self.env[key] = value
@@ -196,6 +209,10 @@ def start(self) -> Self:
         )
 
         logger.info("Container started: %s", self._container.short_id)
+
+        for t in self._transferrables:
+            self._transfer_into_container(t.source, t.destination_in_container, t.mode)
+
         return self
 
     def stop(self, force: bool = True, delete_volume: bool = True) -> None:
@@ -273,6 +290,43 @@ def _configure(self) -> None:
         # placeholder if subclasses want to define this and use the default start method
         pass
 
+    def with_copy_into_container(
+        self, file_content: bytes | PathLike, destination_in_container: str, mode: int = 0o644
+    ):
+        self._transferrables.append(Transferrable(file_content, destination_in_container, mode))
+        return self
+
+    def copy_into_container(self, file_content: bytes, destination_in_container: str, mode: int = 0o644):
+        return self._transfer_into_container(file_content, destination_in_container, mode)
+
+    def _transfer_into_container(self, source: bytes | PathLike, destination_in_container: str, mode: int):
+        if isinstance(source, bytes):
+            file_content = source
+        elif isinstance(source, PathLike):
+            p = pathlib.Path(source)
+            file_content = p.read_bytes()
+        else:
+            raise TypeError("source must be bytes or PathLike")
+
+        fileobj = io.BytesIO()
+        with tarfile.open(fileobj=fileobj, mode="w") as tar:
+            tarinfo = tarfile.TarInfo(name=destination_in_container)
+            tarinfo.size = len(file_content)
+            tarinfo.mode = mode
+            tar.addfile(tarinfo, io.BytesIO(file_content))
+        fileobj.seek(0)
+        rv = self._container.put_archive(path="/", data=fileobj.getvalue())
+        assert rv is True
+
+    def copy_from_container(self, source_in_container: str, destination_on_host: PathLike):
+        tar_stream, _ = self._container.get_archive(source_in_container)
+
+        for chunk in tar_stream:
+            with tarfile.open(fileobj=io.BytesIO(chunk)) as tar:
+                for member in tar.getmembers():
+                    with open(destination_on_host, "wb") as f:
+                        f.write(tar.extractfile(member).read())
+
 
 class Reaper:
     """

core/tests/test_core.py

@@ -1,7 +1,7 @@
 import tempfile
 from pathlib import Path
 
-from testcontainers.core.container import DockerContainer
+from testcontainers.core.container import DockerContainer, Transferrable
 
 
 def test_garbage_collection_is_defensive():
@@ -46,3 +46,121 @@ def test_docker_container_with_env_file():
             assert "[email protected]" in output
             assert "ROOT_URL=example.org/app" in output
             print(output)
+
+
+def test_copy_file_into_container_at_runtime(tmp_path: Path):
+    # Given
+    my_file = tmp_path / "my_file"
+    my_file.write_text("hello world")
+    destination_in_container = "/tmp/my_file"
+
+    with DockerContainer("bash", command="sleep infinity") as container:
+        # When
+        container.copy_into_container(my_file, destination_in_container)
+        result = container.exec(f"cat {destination_in_container}")
+
+    # Then
+    assert result.exit_code == 0
+    assert result.output == b"hello world"
+
+
+def test_copy_file_into_container_at_startup(tmp_path: Path):
+    # Given
+    my_file = tmp_path / "my_file"
+    my_file.write_text("hello world")
+    destination_in_container = "/tmp/my_file"
+
+    container = DockerContainer("bash", command="sleep infinity")
+    container.with_copy_into_container(my_file, destination_in_container)
+
+    with container:
+        # When
+        result = container.exec(f"cat {destination_in_container}")
+
+    # Then
+    assert result.exit_code == 0
+    assert result.output == b"hello world"
+
+
+def test_copy_file_into_container_via_initializer(tmp_path: Path):
+    # Given
+    my_file = tmp_path / "my_file"
+    my_file.write_text("hello world")
+    destination_in_container = "/tmp/my_file"
+
+    with DockerContainer(
+        "bash", command="sleep infinity", transferrables=(Transferrable(my_file, destination_in_container),)
+    ) as container:
+        # When
+        result = container.exec(f"cat {destination_in_container}")
+
+    # Then
+    assert result.exit_code == 0
+    assert result.output == b"hello world"
+
+
+def test_copy_bytes_to_container_at_runtime():
+    # Given
+    file_content = b"hello world"
+    destination_in_container = "/tmp/my_file"
+
+    with DockerContainer("bash", command="sleep infinity") as container:
+        # When
+        container.copy_into_container(file_content, destination_in_container)
+
+        # Then
+        result = container.exec(f"cat {destination_in_container}")
+
+    assert result.exit_code == 0
+    assert result.output == b"hello world"
+
+
+def test_copy_bytes_to_container_at_startup():
+    # Given
+    file_content = b"hello world"
+    destination_in_container = "/tmp/my_file"
+
+    container = DockerContainer("bash", command="sleep infinity")
+    container.with_copy_into_container(file_content, destination_in_container)
+
+    with container:
+        # When
+        result = container.exec(f"cat {destination_in_container}")
+
+    # Then
+    assert result.exit_code == 0
+    assert result.output == b"hello world"
+
+
+def test_copy_bytes_to_container_via_initializer():
+    # Given
+    file_content = b"hello world"
+    destination_in_container = "/tmp/my_file"
+
+    with DockerContainer(
+        "bash", command="sleep infinity", transferrables=(Transferrable(file_content, destination_in_container),)
+    ) as container:
+        # When
+        result = container.exec(f"cat {destination_in_container}")
+
+    # Then
+    assert result.exit_code == 0
+    assert result.output == b"hello world"
+
+
+def test_copy_file_from_container(tmp_path: Path):
+    # Given
+    file_in_container = "/tmp/foo.txt"
+    destination_on_host = tmp_path / "foo.txt"
+    assert not destination_on_host.is_file()
+
+    with DockerContainer("bash", command="sleep infinity") as container:
+        result = container.exec(f'bash -c "echo -n hello world > {file_in_container}"')
+        assert result.exit_code == 0
+
+        # When
+        container.copy_from_container(file_in_container, destination_on_host)
+
+    # Then
+    assert destination_on_host.is_file()
+    assert destination_on_host.read_text() == "hello world"