Added docstrings and finished TODOs in builder

buildrunner/docker/multiplatform_image_builder.py

@@ -9,7 +9,7 @@
 from multiprocessing import Process
 import os
 from platform import system, machine
-from typing import List
+from typing import Iterator, List, Union
 from python_on_whales import docker
 import python_on_whales
 
@@ -151,9 +151,23 @@ def build_image(self,
                     push: bool = True,
                     path: str = '.',
                     file: str = 'Dockerfile',
-                    tags: List[str] = None,):
+                    tags: List[str] = None,
+    ) -> Union[None, python_on_whales.components.image.cli_wrapper.Image, Iterator[str]]:
         """
         Builds the image for the given platform
+
+        Args:
+            name (str): The name of the image
+            platform (str): The platform to build the image for (e.g. linux/amd64)
+            push (bool, optional): Whether to push the image to the registry. Defaults to True.
+            path (str, optional): The path to the Dockerfile. Defaults to '.'.
+            file (str, optional): The path/name of the Dockerfile (ie. <path>/Dockerfile). Defaults to 'Dockerfile'.
+            tags (List[str], optional): The tags to apply to the image. Defaults to None.
+
+        Returns:
+            A `python_on_whales.Image` if a Docker image is loaded
+            in the daemon after the build (the default behavior when
+            calling `docker.build(...)`). Otherwise, `None`.
         """
         if tags is None:
             tags = ['latest']
@@ -166,7 +180,7 @@ def build_image(self,
         LOGGER.debug(f"Building {tagged_names} for {platform}")
 
         # Build the image with the specified tags
-        docker.buildx.build(path, tags=tagged_names, platforms=[platform], push=push, file=file)
+        image = docker.buildx.build(path, tags=tagged_names, platforms=[platform], push=push, file=file)
 
         # Check that the images were built and in the registry
         # Docker search is not currently implemented in python-on-wheels
@@ -178,17 +192,31 @@ def build_image(self,
             except python_on_whales.exceptions.DockerException as err:
                 LOGGER.error(f"Failed to build {tag_name}: {err}")
                 raise err
+        return image
 
     def build(self,
               platforms: List[str],
               path: str = '.',
               file: str = 'Dockerfile',
               name: str = None,
               tags: List[str] = None,
-              do_multiprocessing: bool = False,
-              push=True) -> None:
+              push=True,
+              do_multiprocessing: bool = False) -> List[ImageInfo]:
         """
         Builds the images for the given platforms
+
+        Args:
+            platforms (List[str]): The platforms to build the image for (e.g. linux/amd64)
+            path (str, optional): The path to the Dockerfile. Defaults to '.'.
+            file (str, optional): The path/name of the Dockerfile (ie. <path>/Dockerfile). Defaults to 'Dockerfile'.
+            name (str, optional): The name of the image. Defaults to None.
+            tags (List[str], optional): The tags to apply to the image. Defaults to None.
+            push (bool, optional): Whether to push the image to the registry. Defaults to True.
+            do_multiprocessing (bool, optional): Whether to use multiprocessing to build the images. Defaults to False.
+
+        Returns:
+            List[ImageInfo]: The list of intermediate built images, these images are ephemeral
+            and will be removed when the builder is garbage collected
         """
         LOGGER.debug(f"Building {name}:{tags} for platforms {platforms} from {file}")
 
@@ -224,9 +252,18 @@ def build(self,
     def push(self, name: str, dest_names: List[str] = None) -> None:
         """
         Pushes the image to the remote registry embedded in dest_names or name if dest_names is None
+
+        Args:
+            name (str): The name of the image to push
+            dest_names (List[str], optional): The names of the images to push to. Defaults to None.
+
+        Raises:
+            TimeoutError: If the image fails to push within the specified timeout and retries
         """
         tagged_names = []
         src_names = []
+
+        # Parameters for timeout and retries
         initial_timeout_seconds = 60
         timeout_step_seconds = 60
         timeout_max_seconds = 600
@@ -236,6 +273,7 @@ def push(self, name: str, dest_names: List[str] = None) -> None:
         # only need get tags for one image, since they should be identical
         assert len(src_images) > 0, f"No images found for {name}"
 
+        # Append the tags to the names prior to pushing
         if dest_names is None:
             dest_names = name
             for tag in src_images[0].tags:
@@ -264,25 +302,32 @@ def push(self, name: str, dest_names: List[str] = None) -> None:
                                        " and {timeout_seconds} seconds each try")
             else:
                 # Process finished within timeout
+                LOGGER.info(f"Successfully created multiplatform images {dest_names}")
                 break
             timeout_seconds += timeout_step_seconds
 
             # Cap timeout at max timeout
             timeout_seconds = min(timeout_seconds, timeout_max_seconds)
         return tagged_names
 
-    def _find_native_platform(self, name: str) -> str:
+    def _find_native_platform_images(self, name: str) -> str:
         """
-        Returns the native platform of the host
+        Returns the built native platform image(s) for the given name
+
+        Args:
+            name (str): The name of the image to find
+
+        Returns:
+            str: The name of the native platform image
         """
         host_os = system()
         host_arch = machine()
         LOGGER.debug(f"Finding native platform for {name} for {host_os}/{host_arch}")
         pattern = f'{host_os}-{host_arch}'
 
+        # No images built for this name
         if name not in self._intermediate_built_images.keys() \
            or self._intermediate_built_images[name] == []:  # pylint: disable=consider-iterating-dictionary
-            # TODO Test this case # pylint: disable=fixme
             return None
 
         match_platform = [image for image in self._intermediate_built_images[name] if image.repo.endswith(pattern)]
@@ -305,19 +350,25 @@ def tag_single_platform(self, name: str, tags: List[str] = None,  dest_name: str
         """
         Loads a single platform image into the host registry. If a matching image to the native platform
         is not found, then the first image is loaded.
+
+        Args:
+            name (str): The name of the image to load
+            tags (List[str], optional): The tags to load. Defaults to 'latest'.
+            dest_name (str, optional): The name to load the image as. Defaults to the 'name' arg.
         """
-        # Handles pylint dangerous-default-value error
+        # This is to handle pylint's "dangerous-default-value" error
         if tags is None:
             tags = ['latest']
         LOGGER.debug(f'Tagging {name} with tags {tags} - Dest name: {dest_name}')
-        source_image = self._find_native_platform(name)
+        source_image = self._find_native_platform_images(name)
         if dest_name is None:
             dest_name = name
 
         if self._tagged_images_names.get(name) is None:
             self._tagged_images_names[name] = []
 
-        # TODO - optimize with a single tag command? # pylint: disable=fixme
+        # Tags all the source images with the dest_name and tags
+        # Then removes the intermediate source images
         for image in source_image.formatted_list():
             try:
                 docker.pull(image)

tests/test_multiplatform.py

@@ -11,14 +11,6 @@
 
 TEST_DIR = os.path.basename(os.path.dirname(__file__))
 
-# Best efforts to made sure all images are removed, sometimes they are not.
-# This is a safety net to make sure that docker images built on-demand.
-# The build machines are managed by github, so this is not a concern for us.
-# MOCK_DOCKER = True
-
-# TODO make sure all created created/images are removed after tests
-# TODO check to make sure a context manager is used
-
 def actual_images_match_expected(actual_images, expected_images) -> List[str]:
     missing_images = []
     found = False
@@ -129,7 +121,7 @@ def test_find_native_platform(mock_os,
     mock_arch.return_value = in_mock_arch
     with MultiplatformImageBuilder(use_local_registry=False) as mp:
         mp._intermediate_built_images = built_images
-        found_platform = mp._find_native_platform(name)
+        found_platform = mp._find_native_platform_images(name)
         assert str(found_platform) == str(expected_image)
 
 @pytest.mark.parametrize("name, platforms, expected_image_names",[
@@ -195,7 +187,7 @@ def test_tag_single_platform_multiple_tags(name, platforms, expected_image_names
         found_image = docker.image.list(filters={'reference': f'{built_images[0]}*'})
         assert len(found_image) == 0
 
-    # Check that the image has be removed for host registry
+    # Check that the tagged image has be removed for host registry
     found_image = docker.image.list(filters={'reference': f'{name}*'})
     assert len(found_image) == 0
 
@@ -274,6 +266,48 @@ def test_push():
         print(f'Cleaning up {build_name}')
         docker.image.remove(build_name, force=True)
 
+def test_push_with_dest_names():
+    dest_names = None
+    built_images = None
+    try:
+        with MultiplatformImageBuilder() as remote_mp:
+            reg_add = remote_mp.registry_address()
+            assert reg_add is not None
+
+            tags=['latest', '0.1.0']
+            build_name = 'test-push-image-2001'
+            dest_names = [f'{reg_add}/{build_name}',f'{reg_add}/another-name']
+            platforms = ['linux/arm64','linux/amd64']
+
+            test_path = f'{TEST_DIR}/test-files/multiplatform'
+            with MultiplatformImageBuilder() as mp:
+                built_images = mp.build(name=build_name,
+                                        platforms=platforms,
+                                        path=test_path,
+                                        file=f'{test_path}/Dockerfile',
+                                        do_multiprocessing=False,
+                                        tags=tags)
+
+                assert built_images is not None
+                mp.push(name=build_name, dest_names=dest_names)
+
+                # Make sure the image isn't in the local registry
+                for dest_name in dest_names:
+                    docker.image.remove(dest_name, force=True)
+
+                    # Pull the image from the remote registry to make sure it is there
+                    try:
+                        docker.image.pull(dest_name)
+                    except DockerException as err:
+                        assert False, f'Failed to find/pull {dest_name} from remote registry: {err}'
+                    found_image = docker.image.list(filters={'reference': f'{dest_name}'})
+                    assert len(found_image) == 1
+
+    finally:
+        for dest_name in dest_names:
+            print(f'Cleaning up {dest_name}')
+            docker.image.remove(dest_name, force=True)
+
 @pytest.mark.parametrize("name, platforms, expected_image_names",[
     ('test-build-image-2000',
      ['linux/arm64'],
@@ -365,6 +399,10 @@ def test_build_with_tags(name, tags, platforms, expected_image_names):
         missing_images = actual_images_match_expected(built_images, expected_image_names)
         assert missing_images == [], f'Failed to find {missing_images} in {[image.repo for image in built_images]}'
 
-@pytest.mark.skip
-def test_created_images_are_removed():
-    pass
+def test_no_images_built():
+    """
+    Check that None is returned when no images are built
+    """
+    with MultiplatformImageBuilder() as mp:
+        image = mp._find_native_platform_images('bogus-image-name')
+        assert image is None