Skip to content

Commit 88cf9a8

Browse files
committed
feat(image): introduce DockerImage class for flexible image handling
Redefines the approach to building and managing Docker images within testcontainers by introducing the DockerImage class. This class encapsulates the logic for building images from Dockerfiles and pulling images from repositories. Key Changes: - Implements DockerImage as a central class for image operations, including build, pull, get, and remove. - Add this class as acceptable image param type for DockerContainer and DockerClient - Enables direct Dockerfile support while preserving the option to pull existing images, facilitating a more dynamic testing setup. This refactor addresses feedback on the initial implementation, proposing a cleaner, more extensible design.
1 parent 2db8e6d commit 88cf9a8

File tree

6 files changed

+167
-5
lines changed

6 files changed

+167
-5
lines changed

core/README.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,4 @@ testcontainers-core
44
:code:`testcontainers-core` is the core functionality for spinning up Docker containers in test environments.
55

66
.. autoclass:: testcontainers.core.container.DockerContainer
7+
.. autoclass:: testcontainers.core.image.DockerImage

core/testcontainers/core/container.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
import contextlib
22
from platform import system
3-
from typing import Optional
3+
from typing import Optional, Union
44

55
from docker.models.containers import Container
66

77
from testcontainers.core.docker_client import DockerClient
88
from testcontainers.core.exceptions import ContainerStartException
9+
from testcontainers.core.image import DockerImage
910
from testcontainers.core.utils import inside_container, is_arm, setup_logger
1011
from testcontainers.core.waiting_utils import wait_container_is_ready
1112

@@ -25,11 +26,11 @@ class DockerContainer:
2526
... delay = wait_for_logs(container, "Hello from Docker!")
2627
"""
2728

28-
def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs) -> None:
29+
def __init__(self, image: Union[DockerImage, str], docker_client_kw: Optional[dict] = None, **kwargs) -> None:
2930
self.env = {}
3031
self.ports = {}
3132
self.volumes = {}
32-
self.image = image
33+
self.image = image.get_wrapped_image() if isinstance(image, DockerImage) else image
3334
self._docker = DockerClient(**(docker_client_kw or {}))
3435
self._container = None
3536
self._command = None

core/testcontainers/core/docker_client.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
from docker.errors import NotFound
2323
from docker.models.containers import Container, ContainerCollection
2424

25+
from testcontainers.core.image import DockerImage
26+
2527
from .utils import default_gateway_ip, inside_container, setup_logger
2628

2729
LOGGER = setup_logger(__name__)
@@ -46,7 +48,7 @@ def __init__(self, **kwargs) -> None:
4648
@ft.wraps(ContainerCollection.run)
4749
def run(
4850
self,
49-
image: str,
51+
image: Union[DockerImage, str],
5052
command: Optional[Union[str, list[str]]] = None,
5153
environment: Optional[dict] = None,
5254
ports: Optional[dict] = None,

core/testcontainers/core/image.py

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
import functools as ft
2+
from typing import Optional
3+
4+
import docker
5+
from docker.client import DockerClient
6+
from docker.models.images import Image, ImageCollection
7+
8+
from .utils import setup_logger
9+
10+
LOGGER = setup_logger(__name__)
11+
12+
13+
class DockerImage:
14+
"""
15+
Basic class to manage docker images.
16+
17+
.. doctest::
18+
>>> from testcontainers.core.image import DockerImage
19+
20+
>>> image = DockerImage().from_dockerfile(path="core/tests/", tag="testcontainers/test-image")
21+
>>> image.exists("testcontainers/test-image")
22+
True
23+
>>> image.get("testcontainers/test-image").id
24+
'sha256:...'
25+
>>> image.remove(force=True)
26+
>>> image.exists("testcontainers/test-image")
27+
False
28+
"""
29+
30+
def __init__(self, docker_client_kw: Optional[dict] = None, **kwargs) -> None:
31+
self._docker = DockerClient().from_env(**(docker_client_kw or {}))
32+
33+
def from_dockerfile(self, path: str, tag: str = "local/image") -> "DockerImage":
34+
"""
35+
Build an image from a Dockerfile.
36+
37+
Args:
38+
path (str): Path to the Dockerfile
39+
tag (str): Tag for the image
40+
41+
Returns:
42+
DockerImage: The current instance
43+
"""
44+
self.build(path=path, tag=tag)
45+
return self
46+
47+
def from_image(self, repository: str, tag: str = "latest") -> "DockerImage":
48+
"""
49+
Pull an image from the registry.
50+
51+
Args:
52+
repository (str): Image repository
53+
tag (str): Image tag
54+
55+
Returns:
56+
DockerImage: The current instance
57+
"""
58+
self.pull(repository=repository, tag=tag)
59+
return self
60+
61+
@ft.wraps(ImageCollection.build)
62+
def build(self, **kwargs) -> "DockerImage":
63+
LOGGER.info("Building image from Dockerfile")
64+
self._image, _ = self._docker.images.build(**kwargs)
65+
return self
66+
67+
@ft.wraps(ImageCollection.pull)
68+
def pull(self, **kwargs) -> Image:
69+
LOGGER.info("Pulling image")
70+
self._image = self._docker.images.pull(**kwargs)
71+
return self
72+
73+
@ft.wraps(ImageCollection.get)
74+
def get(self, image: str) -> Image:
75+
LOGGER.info(f"Getting image {image}")
76+
image_obj = self._docker.images.get(image)
77+
return image_obj
78+
79+
@ft.wraps(ImageCollection.remove)
80+
def remove(self, **kwargs) -> None:
81+
LOGGER.info(f"Removing image {self._image}")
82+
self._image.remove(**kwargs)
83+
84+
@property
85+
def id(self) -> str:
86+
return self._image.id
87+
88+
@property
89+
def short_id(self) -> str:
90+
return self._image.short_id
91+
92+
@property
93+
def tags(self) -> dict:
94+
return self._image.tags
95+
96+
def get_wrapped_image(self) -> Image:
97+
return self._image
98+
99+
def exists(self, image: str) -> bool:
100+
"""
101+
Check if the image exists in the local registry.
102+
103+
Args:
104+
image (str): Image name
105+
106+
Returns:
107+
bool: True if the image exists, False otherwise
108+
Raises:
109+
docker.errors.ImageNotFound: If the image does not exist
110+
"""
111+
try:
112+
self.get(image)
113+
return True
114+
except docker.errors.ImageNotFound:
115+
return False

core/tests/test_core.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import pytest
2-
32
from testcontainers.core.container import DockerContainer
3+
from testcontainers.core.image import DockerImage
44
from testcontainers.core.waiting_utils import wait_for_logs
55

66

@@ -29,3 +29,12 @@ def test_can_get_logs():
2929
wait_for_logs(container, "Hello from Docker!")
3030
stdout, stderr = container.get_logs()
3131
assert stdout, "There should be something on stdout"
32+
33+
34+
def test_create_container_from_image():
35+
image = DockerImage().from_dockerfile(path="core/tests/", tag="testcontainers/test-image")
36+
37+
container = DockerContainer(image)
38+
container.start()
39+
container.stop(force=True, delete_volume=True)
40+
image.remove(force=True)

core/tests/test_image.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
import pytest
2+
from testcontainers.core.image import DockerImage
3+
4+
5+
def test_docker_image_from_dockerfile():
6+
image = DockerImage().from_dockerfile(path="core/tests/", tag="testcontainers/test-image")
7+
8+
assert image.exists("testcontainers/test-image") == True
9+
10+
retrieved_image = image.get("testcontainers/test-image")
11+
12+
assert retrieved_image.id == image.id
13+
assert retrieved_image.short_id == image.short_id
14+
assert retrieved_image.tags == image.tags
15+
16+
image.remove(force=True)
17+
18+
assert image.exists("testcontainers/test-image") == False
19+
20+
21+
def test_docker_image_from_image():
22+
image = DockerImage().from_image(repository="alpine")
23+
24+
assert image.exists("alpine") == True
25+
26+
retrieved_image = image.get("alpine")
27+
28+
assert retrieved_image.id == image.id
29+
assert retrieved_image.short_id == image.short_id
30+
assert retrieved_image.tags == image.tags
31+
32+
image.remove(force=True)
33+
34+
assert image.exists("alpine") == False

0 commit comments

Comments
 (0)