chore: clean image transfer and add tests

Signed-off-by: ThibaultFy <[email protected]>

backend/dev-requirements.txt

@@ -21,3 +21,4 @@ mypy==1.4.1
 djangorestframework-stubs==1.8.0
 django-stubs==1.14.0
 celery-types==0.14.0
+docker==6.1.3

backend/image_transfer/__init__.py

@@ -1,6 +1,7 @@
-from image_transfer.decoder import BlobNotFound
-from image_transfer.decoder import ManifestNotFound
+# The image transfer module is a copy of docker charon -> https://github.com/gabrieldemarmiesse/docker-charon
+# Some unused features have been remove and can be found in the original repository.
+
 from image_transfer.decoder import push_payload
 from image_transfer.encoder import make_payload
 
-__all__ = (BlobNotFound, ManifestNotFound, push_payload, make_payload)
+__all__ = (push_payload, make_payload)

backend/image_transfer/common.py

@@ -2,7 +2,6 @@
 
 import json
 from enum import Enum
-from pathlib import Path
 from typing import IO
 from typing import Dict
 from typing import Iterator
@@ -14,6 +13,8 @@
 from dxf import DXFBase
 from pydantic import BaseModel
 
+from image_transfer.exceptions import ManifestContentError
+
 
 class PayloadSide(Enum):
     ENCODER = "ENCODER"
@@ -39,11 +40,13 @@ def __init__(
         dxf_base: DXFBase,
         docker_image_name: str,
         payload_side: PayloadSide,
+        platform: Optional[str] = None,
         content: Optional[str] = None,
     ):
         self.dxf_base = dxf_base
         self.docker_image_name = docker_image_name
         self.payload_side = payload_side
+        self.platform = platform
         self._content = content
 
     @property
@@ -62,11 +65,18 @@ def content(self) -> str:
                     "This makes no sense to fetch the manifest from " "the registry if you're decoding the zip"
                 )
             dxf = DXF.from_base(self.dxf_base, self.repository)
-            self._content = dxf.get_manifest(self.tag)
+            self._content = dxf.get_manifest(self.tag, platform=self.platform)
         return self._content
 
     def get_list_of_blobs(self) -> list[Blob]:
-        manifest_dict = json.loads(self.content)
+        try:
+            manifest_dict = json.loads(self.content)
+        except TypeError:
+            raise ManifestContentError(
+                "The Manifest content must be str, bytes or bytearray. "
+                "Is there several platform available in the manifest ? "
+                "If yes, please specify it."
+            )
         result: list[Blob] = [Blob(self.dxf_base, manifest_dict["config"]["digest"], self.repository)]
         for layer in manifest_dict["layers"]:
             result.append(Blob(self.dxf_base, layer["digest"], self.repository))
@@ -122,9 +132,6 @@ def file_to_generator(file_like: IO) -> Iterator[bytes]:
         yield chunk
 
 
-PROJECT_ROOT = Path(__file__).parents[1]
-
-
 def get_repo_and_tag(docker_image_name: str) -> (str, str):
     return docker_image_name.split(":", 1)
 

backend/image_transfer/decoder.py

@@ -23,14 +23,7 @@
 from image_transfer.common import file_to_generator
 from image_transfer.common import get_repo_and_tag
 from image_transfer.common import progress_as_string
-
-
-class ManifestNotFound(Exception):
-    pass
-
-
-class BlobNotFound(Exception):
-    pass
+from image_transfer.exceptions import ManifestNotFoundError
 
 
 def push_payload(
@@ -45,7 +38,7 @@ def push_payload(
 
     It will iterate over the docker images and push the blobs and the manifests.
 
-    # Arguments
+    Args:
         zip_file: the zip file containing the payload. It can be a `pathlib.Path`, a `str`
             or a file-like object.
         strict: `False` by default. If True, it will raise an error if the
@@ -61,12 +54,12 @@ def push_payload(
         password: the password to use to connect to the registry. Optional
             if the registry does not require authentication.
 
-    # Returns
+    Returns:
         The list of docker images loaded in the registry
         It also includes the list of docker images that were already present
         in the registry and were not included in the payload to optimize the size.
         In other words, it's the argument `docker_images_to_transfer` that you passed
-        to the function `docker_charon.make_payload(...)`.
+        to the function `image_transfer.make_payload(...)`.
     """
     authenticator = Authenticator(username, password)
 
@@ -129,11 +122,11 @@ def check_if_the_docker_image_is_in_the_registry(dxf_base: DXFBase, docker_image
             raise
         error_message = (
             f"The docker image {docker_image} is not present in the "
-            f"registry. But when making the payload, it was specified in "
-            f"`docker_images_already_transferred`."
+            "registry. But when making the payload, it was specified in "
+            "`docker_images_already_transferred`."
         )
         if strict:
-            raise ManifestNotFound(
+            raise ManifestNotFoundError(
                 f"{error_message}\n" f"If you still want to unpack your payload, set `strict=False`."
             )
         else:
@@ -162,4 +155,6 @@ def load_zip_images_in_registry(dxf_base: DXFBase, zip_file: ZipFile, strict: bo
 
 
 def get_payload_descriptor(zip_file: ZipFile) -> PayloadDescriptor:
+    # Replace for Pydantic v2:
+    # PayloadDescriptor.model_validate_json( zip_file.read("payload_descriptor.json").decode())
     return PayloadDescriptor.parse_raw(zip_file.read("payload_descriptor.json").decode())

backend/image_transfer/encoder.py

@@ -74,65 +74,49 @@ def get_blob_with_same_digest(list_of_blobs: list[Blob], digest: str) -> Optiona
             return blob
 
 
-def get_manifest_and_list_of_blobs_to_pull(dxf_base: DXFBase, docker_image: str) -> tuple[Manifest, list[Blob]]:
-    manifest = Manifest(dxf_base, docker_image, PayloadSide.ENCODER)
+def get_manifest_and_list_of_blobs_to_pull(
+    dxf_base: DXFBase,
+    docker_image: str,
+    platform: Optional[str] = None,
+) -> tuple[Manifest, list[Blob]]:
+    manifest = Manifest(dxf_base, docker_image, PayloadSide.ENCODER, platform=platform)
     return manifest, manifest.get_list_of_blobs()
 
 
 def get_manifests_and_list_of_all_blobs(
-    dxf_base: DXFBase, docker_images: Iterator[str]
+    dxf_base: DXFBase, docker_images: Iterator[str], platform: Optional[str] = None
 ) -> tuple[list[Manifest], list[Blob]]:
     manifests = []
     blobs_to_pull = []
     for docker_image in docker_images:
-        manifest, blobs = get_manifest_and_list_of_blobs_to_pull(dxf_base, docker_image)
+        manifest, blobs = get_manifest_and_list_of_blobs_to_pull(dxf_base, docker_image, platform)
         manifests.append(manifest)
         blobs_to_pull += blobs
     return manifests, blobs_to_pull
 
 
-def uniquify_blobs(blobs: list[Blob]) -> list[Blob]:
-    result = []
-    for blob in blobs:
-        if blob.digest not in [x.digest for x in result]:
-            result.append(blob)
-    return result
-
-
-def separate_images_to_transfer_and_images_to_skip(
-    docker_images_to_transfer: list[str], docker_images_already_transferred: list[str]
-) -> tuple[list[str], list[str]]:
-    docker_images_to_transfer_with_blobs = []
-    docker_images_to_skip = []
-    for docker_image in docker_images_to_transfer:
-        if docker_image not in docker_images_already_transferred:
-            docker_images_to_transfer_with_blobs.append(docker_image)
-        else:
-            print(
-                f"Skipping {docker_image} as it has already been transferred",
-                file=sys.stderr,
-            )
-            docker_images_to_skip.append(docker_image)
-    return docker_images_to_transfer_with_blobs, docker_images_to_skip
-
-
 def create_zip_from_docker_images(
     dxf_base: DXFBase,
     docker_images_to_transfer: list[str],
     docker_images_already_transferred: list[str],
     zip_file: ZipFile,
+    platform: Optional[str] = None,
 ) -> None:
     payload_descriptor = PayloadDescriptor.from_images(docker_images_to_transfer, docker_images_already_transferred)
 
     manifests, blobs_to_pull = get_manifests_and_list_of_all_blobs(
-        dxf_base, payload_descriptor.get_images_not_transferred_yet()
+        dxf_base, payload_descriptor.get_images_not_transferred_yet(), platform=platform
+    )
+    _, blobs_already_transferred = get_manifests_and_list_of_all_blobs(
+        dxf_base, docker_images_already_transferred, platform=platform
     )
-    _, blobs_already_transferred = get_manifests_and_list_of_all_blobs(dxf_base, docker_images_already_transferred)
     payload_descriptor.blobs_paths = add_blobs_to_zip(dxf_base, zip_file, blobs_to_pull, blobs_already_transferred)
     for manifest in manifests:
         dest = payload_descriptor.manifests_paths[manifest.docker_image_name]
         zip_file.writestr(dest, manifest.content)
 
+    # Replace for Pydantic v2:
+    # zip_file.writestr("payload_descriptor.json", payload_descriptor.model_dump_json(indent=4))
     zip_file.writestr("payload_descriptor.json", payload_descriptor.json(indent=4))
 
 
@@ -142,17 +126,15 @@ def make_payload(
     docker_images_already_transferred: Optional[list[str]] = None,
     registry: str = "registry-1.docker.io",
     secure: bool = True,
+    platform: Optional[str] = None,
     username: Optional[str] = None,
     password: Optional[str] = None,
 ) -> None:
     """
     Creates a payload from a list of docker images
     All the docker images must be in the same registry.
-    This is currently a limitation of the docker-charon package.
 
-    If you are interested in multi-registries, please open an issue.
-
-    # Arguments
+    Args:
         zip_file: The path to the zip file to create. It can be a `pathlib.Path` or
             a `str`. It's also possible to pass a file-like object. The payload with
             all the docker images is a single zip file.
@@ -164,13 +146,15 @@ def make_payload(
         registry: the registry to push to. It defaults to `registry-1.docker.io` (dockerhub).
         secure: Set to `False` if the registry doesn't support HTTPS (TLS). Default
             is `True`.
+        platform: In case of multi platform images, you can precise which one you want to pull.
         username: The username to use for authentication to the registry. Optional if
             the registry doesn't require authentication.
         password: The password to use for authentication to the registry. Optional if
             the registry doesn't require authentication.
     """
     if docker_images_already_transferred is None:
         docker_images_already_transferred = []
+
     authenticator = Authenticator(username, password)
 
     with DXFBase(host=registry, auth=authenticator.auth, insecure=not secure) as dxf_base:
@@ -180,4 +164,5 @@ def make_payload(
                 docker_images_to_transfer,
                 docker_images_already_transferred,
                 zip_file_opened,
+                platform=platform,
             )

backend/image_transfer/exceptions.py

@@ -0,0 +1,12 @@
+class ManifestNotFoundError(Exception):
+    """The given docker image is not present in the
+    registry. while it was specified in
+    `docker_images_already_transferred`."""
+
+    pass
+
+
+class ManifestContentError(Exception):
+    """The Manifest content must be str, bytes or bytearray."""
+
+    pass

backend/image_transfer/tests/Dockerfile

@@ -0,0 +1,3 @@
+FROM ubuntu:mantic-20230926
+
+RUN echo "hello-world" > /hello-world.txt

backend/image_transfer/tests/conftest.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+
+import docker
+import pytest
+
+
+@pytest.fixture(scope="session")
+def ubuntu_base_image():
+    return "ubuntu:mantic-20230926"
+
+
+@pytest.fixture(scope="session")
+def tag_ubuntu_custom_image():
+    return "ubuntu:augmented"
+
+
+@pytest.fixture(scope="session")
+def base_registry_local_port():
+    return 5555
+
+
+@pytest.fixture(scope="session")
+def docker_client():
+    return docker.from_env()
+
+
+def transfer_to_base_registry(image_name, docker_client, base_registry_local_port):
+    # we transfer the image to the local registry
+    image = docker_client.images.pull(image_name)
+    new_name = f"localhost:{base_registry_local_port}/{image_name}"
+    image.tag(new_name)
+    docker_client.images.push(new_name)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def initialize_local_registry(docker_client, ubuntu_base_image, tag_ubuntu_custom_image, base_registry_local_port):
+    # we create a local registry and add a docker image to it
+    base_registry = docker_client.containers.run(
+        "registry:2",
+        detach=True,
+        ports={"5000/tcp": base_registry_local_port},
+        name="image-transfer-test-registry",
+    )
+    transfer_to_base_registry(ubuntu_base_image, docker_client, base_registry_local_port)
+    # transfer_to_base_registry("busybox:1.36.1", docker_client, base_registry_local_port)
+
+    image, _ = docker_client.images.build(
+        path=str(Path(__file__).parent),
+        tag=f"localhost:{base_registry_local_port}/{tag_ubuntu_custom_image}",
+    )
+    docker_client.images.push(image.tags[0])
+
+    yield
+    base_registry.remove(force=True, v=True)