diff --git a/core/README.rst b/core/README.rst index 2256bd204..94a554fe3 100644 --- a/core/README.rst +++ b/core/README.rst @@ -4,3 +4,4 @@ testcontainers-core :code:`testcontainers-core` is the core functionality for spinning up Docker containers in test environments. .. autoclass:: testcontainers.core.container.DockerContainer +.. autoclass:: testcontainers.core.image.DockerImage diff --git a/core/testcontainers/core/container.py b/core/testcontainers/core/container.py index 6ecc384bd..6d8d82510 100644 --- a/core/testcontainers/core/container.py +++ b/core/testcontainers/core/container.py @@ -1,11 +1,12 @@ import contextlib from platform import system -from typing import Optional +from typing import Optional, Union from docker.models.containers import Container from testcontainers.core.docker_client import DockerClient from testcontainers.core.exceptions import ContainerStartException +from testcontainers.core.image import DockerImage from testcontainers.core.utils import inside_container, is_arm, setup_logger from testcontainers.core.waiting_utils import wait_container_is_ready @@ -25,11 +26,11 @@ class DockerContainer: ... delay = wait_for_logs(container, "Hello from Docker!") """ - def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs) -> None: + def __init__(self, image: Union[DockerImage, str], docker_client_kw: Optional[dict] = None, **kwargs) -> None: self.env = {} self.ports = {} self.volumes = {} - self.image = image + self.image = image.get_wrapped_image() if isinstance(image, DockerImage) else image self._docker = DockerClient(**(docker_client_kw or {})) self._container = None self._command = None diff --git a/core/testcontainers/core/docker_client.py b/core/testcontainers/core/docker_client.py index 05d5377a3..faea2014d 100644 --- a/core/testcontainers/core/docker_client.py +++ b/core/testcontainers/core/docker_client.py @@ -22,6 +22,8 @@ from docker.errors import NotFound from docker.models.containers import Container, ContainerCollection +from testcontainers.core.image import DockerImage + from .utils import default_gateway_ip, inside_container, setup_logger LOGGER = setup_logger(__name__) @@ -46,7 +48,7 @@ def __init__(self, **kwargs) -> None: @ft.wraps(ContainerCollection.run) def run( self, - image: str, + image: Union[DockerImage, str], command: Optional[Union[str, list[str]]] = None, environment: Optional[dict] = None, ports: Optional[dict] = None, diff --git a/core/testcontainers/core/image.py b/core/testcontainers/core/image.py new file mode 100644 index 000000000..9a1f8712d --- /dev/null +++ b/core/testcontainers/core/image.py @@ -0,0 +1,115 @@ +import functools as ft +from typing import Optional + +import docker +from docker.client import DockerClient +from docker.models.images import Image, ImageCollection + +from .utils import setup_logger + +LOGGER = setup_logger(__name__) + + +class DockerImage: + """ + Basic class to manage docker images. + + .. doctest:: + >>> from testcontainers.core.image import DockerImage + + >>> image = DockerImage().from_dockerfile(path="core/tests/", tag="testcontainers/test-image") + >>> image.exists("testcontainers/test-image") + True + >>> image.get("testcontainers/test-image").id + 'sha256:...' + >>> image.remove(force=True) + >>> image.exists("testcontainers/test-image") + False + """ + + def __init__(self, docker_client_kw: Optional[dict] = None, **kwargs) -> None: + self._docker = DockerClient().from_env(**(docker_client_kw or {})) + + def from_dockerfile(self, path: str, tag: str = "local/image") -> "DockerImage": + """ + Build an image from a Dockerfile. + + Args: + path (str): Path to the Dockerfile + tag (str): Tag for the image + + Returns: + DockerImage: The current instance + """ + self.build(path=path, tag=tag) + return self + + def from_image(self, repository: str, tag: str = "latest") -> "DockerImage": + """ + Pull an image from the registry. + + Args: + repository (str): Image repository + tag (str): Image tag + + Returns: + DockerImage: The current instance + """ + self.pull(repository=repository, tag=tag) + return self + + @ft.wraps(ImageCollection.build) + def build(self, **kwargs) -> "DockerImage": + LOGGER.info("Building image from Dockerfile") + self._image, _ = self._docker.images.build(**kwargs) + return self + + @ft.wraps(ImageCollection.pull) + def pull(self, **kwargs) -> Image: + LOGGER.info("Pulling image") + self._image = self._docker.images.pull(**kwargs) + return self + + @ft.wraps(ImageCollection.get) + def get(self, image: str) -> Image: + LOGGER.info(f"Getting image {image}") + image_obj = self._docker.images.get(image) + return image_obj + + @ft.wraps(ImageCollection.remove) + def remove(self, **kwargs) -> None: + LOGGER.info(f"Removing image {self._image}") + self._image.remove(**kwargs) + + @property + def id(self) -> str: + return self._image.id + + @property + def short_id(self) -> str: + return self._image.short_id + + @property + def tags(self) -> dict: + return self._image.tags + + def get_wrapped_image(self) -> Image: + return self._image + + def exists(self, image: str) -> bool: + """ + Check if the image exists in the local registry. + + Args: + image (str): Image name + + Returns: + bool: True if the image exists, False otherwise + Raises: + docker.errors.ImageNotFound: If the image does not exist + """ + try: + self.get(image) + return True + except docker.errors.ImageNotFound: + return False diff --git a/core/tests/test_core.py b/core/tests/test_core.py index a00be1f02..900ae7a5c 100644 --- a/core/tests/test_core.py +++ b/core/tests/test_core.py @@ -1,6 +1,6 @@ import pytest - from testcontainers.core.container import DockerContainer +from testcontainers.core.image import DockerImage from testcontainers.core.waiting_utils import wait_for_logs @@ -29,3 +29,12 @@ def test_can_get_logs(): wait_for_logs(container, "Hello from Docker!") stdout, stderr = container.get_logs() assert stdout, "There should be something on stdout" + + +def test_create_container_from_image(): + image = DockerImage().from_dockerfile(path="core/tests/", tag="testcontainers/test-image") + + container = DockerContainer(image) + container.start() + container.stop(force=True, delete_volume=True) + image.remove(force=True) diff --git a/core/tests/test_image.py b/core/tests/test_image.py new file mode 100644 index 000000000..98dc35221 --- /dev/null +++ b/core/tests/test_image.py @@ -0,0 +1,34 @@ +import pytest +from testcontainers.core.image import DockerImage + + +def test_docker_image_from_dockerfile(): + image = DockerImage().from_dockerfile(path="core/tests/", tag="testcontainers/test-image") + + assert image.exists("testcontainers/test-image") == True + + retrieved_image = image.get("testcontainers/test-image") + + assert retrieved_image.id == image.id + assert retrieved_image.short_id == image.short_id + assert retrieved_image.tags == image.tags + + image.remove(force=True) + + assert image.exists("testcontainers/test-image") == False + + +def test_docker_image_from_image(): + image = DockerImage().from_image(repository="alpine") + + assert image.exists("alpine") == True + + retrieved_image = image.get("alpine") + + assert retrieved_image.id == image.id + assert retrieved_image.short_id == image.short_id + assert retrieved_image.tags == image.tags + + image.remove(force=True) + + assert image.exists("alpine") == False