Image clean is now optional

Author: Tranquility2

core/testcontainers/core/generic.py

@@ -99,8 +99,8 @@ class CustomContainer(DockerContainer):
     :param tag: Tag for the image to be built (default: None)
     """
 
-    def __init__(self, path: str, tag: Optional[str]) -> None:
-        self.docker_image = DockerImage(path=path, tag=tag).build()
+    def __init__(self, path: str, tag: Optional[str], image_cleanup: bool = True) -> None:
+        self.docker_image = DockerImage(path=path, tag=tag, clean_up=image_cleanup).build()
         super().__init__(str(self.docker_image))
 
     def get_url(self) -> str:

core/testcontainers/core/image.py

@@ -30,12 +30,14 @@ def __init__(
         path: str,
         docker_client_kw: Optional[dict] = None,
         tag: Optional[str] = None,
+        clean_up: bool = True,
         **kwargs,
     ) -> None:
         self.tag = tag
         self.path = path
         self.id = None
         self._docker = DockerClient(**(docker_client_kw or {}))
+        self.clean_up = clean_up
         self._kwargs = kwargs
 
     def build(self, **kwargs) -> Self:
@@ -46,7 +48,7 @@ def build(self, **kwargs) -> Self:
         return self
 
     @property
-    def short_id(self):
+    def short_id(self) -> str:
         """
         The ID of the image truncated to 12 characters, without the ``sha256:`` prefix.
         """
@@ -61,7 +63,8 @@ def remove(self, force=True, noprune=False) -> None:
         :param force: Remove the image even if it is in use
         :param noprune: Do not delete untagged parent images
         """
-        if self._image:
+        if self._image and self.clean_up:
+            logger.info(f"Cleanup - Removing image {self.short_id}")
             self._image.remove(force=force, noprune=noprune)
         self.get_docker_client().client.close()
 

core/tests/test_core.py

@@ -1,10 +1,14 @@
 import pytest
 import tempfile
+import random
+
+from typing import Optional
 
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.image import DockerImage
 from testcontainers.core.waiting_utils import wait_for_logs
 from testcontainers.core.generic import CustomContainer
+from testcontainers.core.docker_client import DockerClient
 
 
 def test_timeout_is_raised_when_waiting_for_logs():
@@ -36,41 +40,64 @@ def test_can_get_logs():
         assert stdout, "There should be something on stdout"
 
 
+def check_for_image(image_short_id: str) -> bool:
+    """
+    Check if the image with the given short_id is present in the list of images
+    """
+    client = DockerClient()
+    images = client.client.images.list()
+    return any(image.short_id.endswith(image_short_id) for image in images)
+
+
+@pytest.mark.parametrize("test_cleanup", [True, False])
 @pytest.mark.parametrize("test_image_tag", [None, "test-image:latest"])
-def test_docker_image(test_image_tag):
+def test_docker_image(test_image_tag: Optional[str], test_cleanup: bool):
     with tempfile.TemporaryDirectory() as temp_directory:
+        # It's important to use a random string to avoid image caching
+        random_string = "Hello from Docker Image! " + str(random.randint(0, 1000))
         with open(f"{temp_directory}/Dockerfile", "w") as f:
             f.write(
-                """
+                f"""
                 FROM alpine:latest
-                CMD echo "Hello from Docker Image!"
+                CMD echo "{random_string}"
                 """
             )
             f.close()
-            with DockerImage(path=temp_directory, tag=test_image_tag) as image:
+            with DockerImage(path=temp_directory, tag=test_image_tag, clean_up=test_cleanup) as image:
+                image_short_id = image.short_id
                 assert image.tag is test_image_tag
                 assert image.short_id is not None
                 logs = image.get_logs()
                 assert isinstance(logs, list)
                 assert logs[0] == {"stream": "Step 1/2 : FROM alpine:latest"}
-                assert logs[3] == {"stream": 'Step 2/2 : CMD echo "Hello from Docker Image!"'}
+                assert logs[3] == {"stream": f'Step 2/2 : CMD echo "{random_string}"'}
                 with DockerContainer(str(image)) as container:
-                    wait_for_logs(container, "Hello from Docker Image!")
-                    assert container._container.image.short_id.endswith(image.short_id)
+                    wait_for_logs(container, random_string)
+                    assert container._container.image.short_id.endswith(image_short_id)
 
+            if not test_cleanup:
+                assert check_for_image(image_short_id)
 
+
+@pytest.mark.parametrize("test_image_cleanup", [True, False])
 @pytest.mark.parametrize("test_image_tag", [None, "custom-image:test"])
-def test_custom_container(test_image_tag):
+def test_custom_container(test_image_tag: Optional[str], test_image_cleanup: bool):
     with tempfile.TemporaryDirectory() as temp_directory:
+        # It's important to use a random string to avoid image caching
+        random_string = "This is a custom Docker Image! " + str(random.randint(0, 1000))
         with open(f"{temp_directory}/Dockerfile", "w") as f:
             f.write(
-                """
+                f"""
                 FROM alpine:latest
-                CMD echo "This is a custom Docker Image!"
+                CMD echo "{random_string}"
                 """
             )
             f.close()
-            with CustomContainer(path=temp_directory, tag=test_image_tag) as server:
+            with CustomContainer(path=temp_directory, tag=test_image_tag, image_cleanup=test_image_cleanup) as server:
+                image_short_id = server.docker_image.short_id
                 image_build_logs = server.docker_image.get_logs()
-                assert image_build_logs[3] == {"stream": 'Step 2/2 : CMD echo "This is a custom Docker Image!"'}
-                wait_for_logs(server, "This is a custom Docker Image!")
+                assert image_build_logs[3] == {"stream": f'Step 2/2 : CMD echo "{random_string}"'}
+                wait_for_logs(server, random_string)
+
+            if not test_image_cleanup:
+                assert check_for_image(image_short_id)