From ecb2d6bdd3670cdbddb03c924971d3e8a4adbf49 Mon Sep 17 00:00:00 2001
From: saville <saville@adobe.com>
Date: Thu, 21 Dec 2023 12:10:47 -0700
Subject: [PATCH 1/3] Refactor multiplatform image builder image handling code

Create a class to track the built image and final image ref and tags
The tests are still broken, but will be fixed in a later commit
---
 buildrunner/__init__.py                       |  11 +-
 buildrunner/docker/image_info.py              | 161 ++++++
 .../docker/multiplatform_image_builder.py     | 524 ++++++------------
 buildrunner/steprunner/tasks/__init__.py      |  30 -
 buildrunner/steprunner/tasks/build.py         |  36 +-
 buildrunner/steprunner/tasks/push.py          |  66 ++-
 tests/test_multiplatform.py                   | 396 ++++---------
 7 files changed, 508 insertions(+), 716 deletions(-)
 create mode 100644 buildrunner/docker/image_info.py

diff --git a/buildrunner/__init__.py b/buildrunner/__init__.py
index 36ed3a78..8be3f301 100644
--- a/buildrunner/__init__.py
+++ b/buildrunner/__init__.py
@@ -644,16 +644,11 @@ def run(self):  # pylint: disable=too-many-statements,too-many-branches,too-many
                         "Push requested--pushing generated images/packages to remote registries/repositories\n"
                     )
                     # Push multi-platform images
-                    if multi_platform.tagged_images_names:
+                    if multi_platform.num_built_images:
                         self.log.write(
-                            f"===> Pushing {len(multi_platform.tagged_images_names)} multiplatform image(s)\n"
+                            f"===> Pushing {multi_platform.num_built_images} multiplatform image(s)\n"
                         )
-                        for (
-                            local_name,
-                            dest_name,
-                        ) in multi_platform.tagged_images_names.items():
-                            self.log.write(f"Pushing {local_name} to {dest_name}\n")
-                            multi_platform.push(name=local_name, dest_names=dest_name)
+                        multi_platform.push(self.log)
 
                     # Push single platform images
                     _docker_client = docker.new_client(timeout=self.docker_timeout)
diff --git a/buildrunner/docker/image_info.py b/buildrunner/docker/image_info.py
new file mode 100644
index 00000000..ccdcc88e
--- /dev/null
+++ b/buildrunner/docker/image_info.py
@@ -0,0 +1,161 @@
+"""
+Copyright 2023 Adobe
+All Rights Reserved.
+
+NOTICE: Adobe permits you to use, modify, and distribute this file in accordance
+with the terms of the Adobe license agreement accompanying it.
+"""
+
+import logging
+import platform as python_platform
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+
+LOGGER = logging.getLogger(__name__)
+
+# Maps from lower-case platform.system() result to platform prefix in docker
+PLATFORM_SYSTEM_MAPPINGS = {"darwin": "linux"}
+PLATFORM_MACHINE_MAPPINGS = {
+    "aarch64": "arm",
+    "x86_64": "amd",
+}
+
+
+class TaggedImageInfo(BaseModel):
+    # The image repo with no tag (e.g. repo/image-name)
+    repo: str
+    # All tags for this image
+    tags: List[str]
+
+    @property
+    def image_refs(self) -> List[str]:
+        """Returns a list of image names with tags"""
+        return [f"{self.repo}:{tag}" for tag in self.tags]
+
+    def __str__(self) -> str:
+        return f'{self.repo}:{",".join(self.tags)}'
+
+
+class BuiltTaggedImage(BaseModel):
+    # The image repo with no tag (e.g. repo/image-name)
+    repo: str
+    # The single tag for this image
+    tag: str
+    # Digest for the image
+    digest: str
+    # Platform for the image
+    platform: str
+
+    @property
+    def trunc_digest(self) -> str:
+        return self.digest.replace("sha256:", "")[:12]
+
+    @property
+    def image_ref(self) -> str:
+        """Returns an image name with the tag"""
+        return f"{self.repo}:{self.tag}"
+
+
+class BuiltImageInfo(BaseModel):
+    """
+    Contains information about images created during the build,
+    including the final tagged image name and any tags (once generated).
+    """
+
+    # A unique ID for this built image
+    id: str
+    # The built images by each platform
+    images_by_platform: Dict[str, BuiltTaggedImage] = {}
+    # The images that should be created from the built images referenced in each instance
+    tagged_images: List[TaggedImageInfo] = []
+
+    @property
+    def platforms(self) -> List[str]:
+        return list(self.images_by_platform.keys())
+
+    @property
+    def built_images(self) -> List[BuiltTaggedImage]:
+        return list(self.images_by_platform.values())
+
+    def image_for_platform(self, platform: str) -> Optional[BuiltTaggedImage]:
+        """Retrieves the built image for the given platform."""
+        if platform not in self.images_by_platform:
+            return None
+        return self.images_by_platform[platform]
+
+    def add_platform_image(self, platform: str, image_info: BuiltTaggedImage) -> None:
+        if platform in self.images_by_platform:
+            raise ValueError(
+                f"Cannot add image {image_info} for platform {platform} since it already exists in this built image info"
+            )
+        self.images_by_platform[platform] = image_info
+
+    def add_tagged_image(self, repo: str, tags: List[str]) -> TaggedImageInfo:
+        """
+        Creates a tagged image instance, adds it to the final image tags, then returns it.
+        :arg repo: The image repo without the tag
+        :arg tags: The tags to use when tagging the image
+        """
+        tagged_image = TaggedImageInfo(
+            repo=repo,
+            tags=tags,
+        )
+        self.tagged_images.append(tagged_image)
+        return tagged_image
+
+    @property
+    def native_platform_image(self) -> BuiltTaggedImage:
+        """
+        Returns the built native platform image(s) matching the current machine's platform or the first built
+        image if none matches.
+
+        Args:
+            name (str): The name of the image to find
+
+        Returns:
+            str: The name of the native platform image
+        """
+        # Converts the python platform results to docker-equivalents, and then matches based on the prefix
+        # e.g. Darwin/aarch64 would become linux/arm which would then match linux/arm64/v7, etc
+        native_system = python_platform.system().lower()
+        native_machine = python_platform.machine()
+        native_system = PLATFORM_SYSTEM_MAPPINGS.get(native_system, native_system)
+        native_machine = PLATFORM_MACHINE_MAPPINGS.get(native_machine, native_machine)
+        LOGGER.debug(
+            f"Finding native platform for {self} for {native_system}/{native_machine}"
+        )
+
+        # Search for the image matching the native platform
+        # Uses dash since this is matching on tag names which do not support forward slashes
+        pattern = f"{native_system}-{native_machine}"
+        matched_image = None
+        for image in self.images_by_platform.values():
+            if image.tag.replace(f"{self.id}-", "").startswith(pattern):
+                matched_image = image
+
+        # Still no matches found, get the first image and log a warning
+        if not matched_image:
+            platform = self.platforms[0]
+            LOGGER.warning(
+                f"Could not find image matching native platform {pattern} for {self}, using platform {platform} instead locally"
+            )
+            return self.images_by_platform[platform]
+
+        return matched_image
+
+    def __str__(self):
+        num_platforms = len(self.platforms)
+        if num_platforms == 0:
+            return f"<{self.id}>"
+        elif num_platforms == 1:
+            return f"<{self.images_by_platform[self.platforms[0]]} ({self.id})>"
+        image_strs = [
+            f"{platform}: {image}"
+            for platform, image in self.images_by_platform.items()
+        ]
+        return f"<{', '.join(image_strs)} ({self.id})>"
+
+    def __repr__(self):
+        return self.__str__()
diff --git a/buildrunner/docker/multiplatform_image_builder.py b/buildrunner/docker/multiplatform_image_builder.py
index 4ffc01cc..7899939d 100644
--- a/buildrunner/docker/multiplatform_image_builder.py
+++ b/buildrunner/docker/multiplatform_image_builder.py
@@ -7,11 +7,11 @@
 """
 import logging
 import os
-from multiprocessing import Manager, Process
-from platform import machine, system
+import platform as python_platform
 import shutil
 import tempfile
 import uuid
+from multiprocessing import Process, SimpleQueue
 from typing import Dict, List, Optional, Union
 
 import python_on_whales
@@ -19,82 +19,17 @@
 from python_on_whales import docker
 from retry import retry
 
+from buildrunner.utils import ConsoleLogger
 from buildrunner.docker import get_dockerfile
+from buildrunner.docker.image_info import BuiltImageInfo, BuiltTaggedImage
 
 
 LOGGER = logging.getLogger(__name__)
+OUTPUT_LINE = "-----------------------------------------------------------------"
 
 PUSH_TIMEOUT = 300
 
 
-class ImageInfo:
-    """Image information repo with associated tags"""
-
-    def __init__(
-        self,
-        repo: str,
-        tags: List[str] = None,
-        platform: str = None,
-        digest: str = None,
-    ):
-        """
-        Args:
-            repo (str): The repo name for the image.
-            tags (List[str], optional): The tags for the image. Defaults to None.
-            platform (str, optional): The platform for the image. Defaults to None.
-            digest (str, optional): The digest id for the image. Defaults to None.
-        """
-        self._repo = repo
-
-        if tags is None:
-            self._tags = ["latest"]
-        else:
-            self._tags = tags
-
-        self._platform = platform
-        self._digest = digest
-
-    @property
-    def repo(self) -> str:
-        """The repo name for the image."""
-        return self._repo
-
-    @property
-    def tags(self) -> List[str]:
-        """The tags for the image."""
-        return self._tags
-
-    @property
-    def digest(self) -> str:
-        """The digest id for the image."""
-        return self._digest
-
-    def trunc_digest(self) -> str:
-        """The truncated digest id for the image."""
-        if self._digest is None:
-            return "Digest not available"
-        return self._digest.replace("sha256:", "")[:12]
-
-    @property
-    def platform(self) -> str:
-        """The platform for the image."""
-        return self._platform
-
-    def formatted_list(self) -> List[str]:
-        """Returns a list of formatted image names"""
-        return [f"{self._repo}:{tag}" for tag in self._tags]
-
-    def __str__(self):
-        """Returns a string representation of the image info"""
-        if len(self._tags) == 1:
-            return f"{self._repo}:{self._tags[0]}"
-        return f"{self._repo} tags={self._tags}"
-
-    def __repr__(self):
-        """Returns a string representation of the image info"""
-        return self.__str__()
-
-
 class RegistryInfo:
     """Registry information"""
 
@@ -159,9 +94,7 @@ def __init__(
                 f'for builders {", ".join(cache_builders) if cache_builders else "(all)"}'
             )
 
-        # key is destination image name, value is list of built images
-        self._intermediate_built_images = {}
-        self._tagged_images_names = {}
+        self._built_images: List[BuiltImageInfo] = []
         self._local_registry_is_running = False
 
     def __enter__(self):
@@ -171,45 +104,11 @@ def __exit__(self, exc_type, exc_value, traceback):
         if self._local_registry_is_running:
             self._stop_local_registry()
 
-        # Removes all intermediate built images
-        if self._intermediate_built_images:
-            for name, images in self._intermediate_built_images.items():
-                for image in images:
-                    for tag in image.tags:
-                        LOGGER.debug(f"Removing image {image.repo}:{tag} for {name}")
-                        docker.image.remove(f"{image.repo}:{tag}", force=True)
-
-        # Removes all tagged images if cleanup_images is set
-        if self._tagged_images_names and self._cleanup_images:
-            for name, images in self._tagged_images_names.items():
-                for image in images:
-                    LOGGER.debug(f"Removing tagged image {image} for {name}")
-                    docker.image.remove(image, force=True)
-
     @property
-    def disable_multi_platform(self) -> int:
-        """Returns the ip address of the local registry"""
+    def disable_multi_platform(self) -> bool:
+        """Returns true if multi-platform builds are disabled by configuration, false otherwise"""
         return self._disable_multi_platform
 
-    @property
-    def registry_ip(self) -> int:
-        """Returns the ip address of the local registry"""
-        return self._mp_registry_info.ip_addr
-
-    @property
-    def registry_port(self) -> int:
-        """Returns the port of the local registry"""
-        return self._mp_registry_info.port
-
-    @property
-    def tagged_images_names(self) -> Dict[str, List[str]]:
-        """Returns a list of all the tagged images names"""
-        return self._tagged_images_names
-
-    def is_multiplatform(self, name: str) -> bool:
-        """Returns True if the image with name is multiplatform"""
-        return len(self._intermediate_built_images.get(name, [])) > 0
-
     def registry_address(self) -> str:
         """Returns the address of the local registry"""
         return f"{self._mp_registry_info.ip_addr}:{self._mp_registry_info.port}"
@@ -281,10 +180,10 @@ def _get_build_cache_options(self, builder: Optional[str]) -> dict:
     def _build_with_inject(
         self,
         inject: dict,
-        tagged_names: List[str],
+        image_ref: str,
         platform: str,
         path: str,
-        file: str,
+        dockerfile: str,
         build_args: dict,
         builder: Optional[str],
         cache: bool = False,
@@ -292,7 +191,7 @@ def _build_with_inject(
     ) -> None:
         if not path or not os.path.isdir(path):
             LOGGER.warning(
-                f"Failed to inject {inject} for {tagged_names} since path {path} isn't a directory."
+                f"Failed to inject {inject} for {image_ref} since path {path} isn't a directory."
             )
             return
 
@@ -325,10 +224,10 @@ def _build_with_inject(
 
             docker.buildx.build(
                 context_dir,
-                tags=tagged_names,
+                tags=[image_ref],
                 platforms=[platform],
                 load=True,
-                file=file,
+                file=dockerfile,
                 builder=builder,
                 build_args=build_args,
                 cache=cache,
@@ -336,6 +235,10 @@ def _build_with_inject(
                 **self._get_build_cache_options(builder),
             )
 
+    @staticmethod
+    def _get_image_digest(image_ref: str) -> str:
+        return docker.buildx.imagetools.inspect(image_ref).config.digest
+
     # pylint: disable=too-many-arguments
     @retry(
         python_on_whales.exceptions.DockerException,
@@ -347,40 +250,37 @@ def _build_with_inject(
     )
     def _build_single_image(
         self,
-        name: str,
+        queue: SimpleQueue,
+        image_ref: str,
         platform: str,
         path: str,
-        file: str,
-        tags: List[str],
+        dockerfile: str,
         build_args: dict,
-        mp_image_name: str,
         inject: dict,
         cache: bool = False,
         pull: bool = False,
     ) -> None:
         """
-        Builds a single image for the given platform
+        Builds a single image for the given platform.
 
         Args:
-            name (str): The name of the image
+            queue (SimpleQueue): The queue to put the message digest on
+            image_ref (str): The image ref for the new image
             platform (str): The platform to build the image for (e.g. linux/amd64)
             path (str): The path to the Dockerfile.
-            file (str): The path/name of the Dockerfile (ie. <path>/Dockerfile).
-            tags (List[str]): The tags to apply to the image.
+            dockerfile (str): The path/name of the Dockerfile (i.e. <path>/Dockerfile).
             build_args (dict): The build args to pass to docker.
-            mp_image_name (str): The multi-platform name of the image.
             inject (dict): The files to inject into the build context.
         """
-        assert os.path.isdir(path) and os.path.exists(f"{file}"), (
-            f"Either path {path}({os.path.isdir(path)}) or file "
-            f"'{file}'({os.path.exists(f'{file}')}) does not exist!"
+        assert os.path.isdir(path) and os.path.exists(dockerfile), (
+            f"Either path {path} ({os.path.isdir(path)}) or file "
+            f"'{dockerfile}' ({os.path.exists(dockerfile)}) does not exist!"
         )
 
-        tagged_names = [f"{name}:{tag}" for tag in tags]
         builder = (
             self._platform_builders.get(platform) if self._platform_builders else None
         )
-        LOGGER.debug(f"Building tagged images {tagged_names}")
+        LOGGER.debug(f"Building image {image_ref} for platform {platform}")
         print(
             f"Building image for platform {platform} with {builder or 'default'} builder"
         )
@@ -388,10 +288,10 @@ def _build_single_image(
         if inject and isinstance(inject, dict):
             self._build_with_inject(
                 inject=inject,
-                tagged_names=tagged_names,
+                image_ref=image_ref,
                 platform=platform,
                 path=path,
-                file=file,
+                dockerfile=dockerfile,
                 build_args=build_args,
                 builder=builder,
                 cache=cache,
@@ -400,10 +300,10 @@ def _build_single_image(
         else:
             docker.buildx.build(
                 path,
-                tags=tagged_names,
+                tags=[image_ref],
                 platforms=[platform],
                 load=True,
-                file=file,
+                file=dockerfile,
                 build_args=build_args,
                 builder=builder,
                 cache=cache,
@@ -411,36 +311,13 @@ def _build_single_image(
                 **self._get_build_cache_options(builder),
             )
         # Push after the initial load to support remote builders that cannot access the local registry
-        docker.push(tagged_names)
+        docker.push([image_ref])
 
-        # Check that the images were built and in the registry
-        # Docker search is not currently implemented in python-on-wheels
-        image_id = None
-        for tag_name in tagged_names:
-            try:
-                images = docker.image.pull(tag_name)
-                assert images, f"Failed to build {tag_name}"
-                image_id = docker.image.inspect(tag_name).id
-                # Removes the image from host, if this fails it is considered a warning
-                try:
-                    LOGGER.debug(f"Removing {tag_name}")
-                    docker.image.remove(tag_name, force=True)
-                except python_on_whales.exceptions.DockerException as err:
-                    LOGGER.warning(f"Failed to remove {images}: {err}")
-            except python_on_whales.exceptions.DockerException as err:
-                LOGGER.error(f"Failed to build {tag_name}: {err}")
-                raise err
-
-        self._intermediate_built_images[mp_image_name].append(
-            ImageInfo(
-                repo=name,
-                tags=tags,
-                platform=platform,
-                digest=image_id,
-            )
-        )
+        # Retrieve the digest and put it on the queue
+        image_digest = self._get_image_digest(image_ref)
+        queue.put((image_ref, image_digest))
 
-    def get_single_platform_to_build(self, platforms: List[str]) -> str:
+    def _get_single_platform_to_build(self, platforms: List[str]) -> str:
         """Returns the platform to build for single platform flag"""
 
         assert isinstance(platforms, list), f"Expected list, but got {type(platforms)}"
@@ -448,9 +325,8 @@ def get_single_platform_to_build(self, platforms: List[str]) -> str:
             len(platforms) > 0
         ), f"Expected at least one platform, but got {len(platforms)}"
 
-        host_system = system()
-        host_machine = machine()
-        native_platform = None
+        host_system = python_platform.system()
+        host_machine = python_platform.machine()
 
         if host_system in ("Darwin", "linux"):
             native_platform = f"linux/{host_machine}"
@@ -469,109 +345,103 @@ def build_multiple_images(
         platforms: List[str],
         path: str = ".",
         file: str = "Dockerfile",
-        mp_image_name: str = None,
-        tags: List[str] = None,
         do_multiprocessing: bool = True,
         build_args: dict = None,
         inject: dict = None,
         cache: bool = False,
         pull: bool = False,
-    ) -> List[ImageInfo]:
+    ) -> BuiltImageInfo:
         """
         Builds multiple images for the given platforms. One image will be built for each platform.
 
-        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".
-            mp_image_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.
-            do_multiprocessing (bool, optional): Whether to use multiprocessing to build the images. Defaults to True.
-            build_args (dict, optional): The build args to pass to docker. Defaults to None.
-            inject (dict, optional): The files to inject into the build context. Defaults to None.
-            cache (bool, optional): If true, enables cache, defaults to False
-            pull (bool, optional): If true, pulls image before build, defaults to False
-            inject (dict, optional): The files to inject into the build context. Defaults to None.
-
-        Returns:
-            List[ImageInfo]: The list of intermediate built images, these images are ephemeral
-            and will be removed when the builder is garbage collected
+        :arg platforms: The platforms to build the image for (e.g. linux/amd64)
+        :arg path: The path to the Dockerfile. Defaults to ".".
+        :arg file: The path/name of the Dockerfile (i.e. <path>/Dockerfile). Defaults to "Dockerfile".
+        :arg do_multiprocessing: Whether to use multiprocessing to build the images. Defaults to True.
+        :arg build_args: The build args to pass to docker. Defaults to None.
+        :arg inject: The files to inject into the build context. Defaults to None.
+        :arg cache: If true, enables cache, defaults to False
+        :arg pull: If true, pulls image before build, defaults to False
+        :return: A BuiltImageInfo instance that describes the built images for each platform and can be used to track
+                 final images as well
         """
-
-        def get_path(file):
-            if os.path.exists(file):
-                return os.path.dirname(file)
-            return "."
-
         if build_args is None:
             build_args = {}
         build_args["DOCKER_REGISTRY"] = self._docker_registry
 
         # It is not valid to pass None for the path when building multi-platform images
         if not path:
-            path = get_path(file)
+            if os.path.exists(file):
+                path = os.path.dirname(file)
+            else:
+                path = "."
 
         dockerfile, cleanup_dockerfile = get_dockerfile(file)
 
+        # Track this newly built image
+        built_image = BuiltImageInfo(id=str(uuid.uuid4()))
+        self._built_images.append(built_image)
+
         LOGGER.debug(
-            f"Building {mp_image_name}:{tags} for platforms {platforms} from {dockerfile}"
+            f"Building multi-platform images {built_image} for platforms {platforms} from {dockerfile}"
         )
 
         if self._use_local_registry and not self._local_registry_is_running:
             # Starts local registry container to do ephemeral image storage
             self._start_local_registry()
 
-        if tags is None:
-            tags = ["latest"]
-
-        # Updates name to be compatible with docker
-        image_prefix = "buildrunner-mp"
-        sanitized_name = f"{image_prefix}-{str(uuid.uuid4())}"
-        base_image_name = f"{self._mp_registry_info.ip_addr}:{self._mp_registry_info.port}/{sanitized_name}"
-
-        # Keeps track of the built images {name: [ImageInfo(image_names)]]}
-        manager = Manager()
-        self._intermediate_built_images[mp_image_name] = manager.list()
-        line = "-----------------------------------------------------------------"
+        # The name should stay consistent and the tags are variable per build step/platform
+        repo = f"{self._mp_registry_info.ip_addr}:{self._mp_registry_info.port}/buildrunner-mp"
 
         if self._disable_multi_platform:
-            platforms = [self.get_single_platform_to_build(platforms)]
+            platforms = [self._get_single_platform_to_build(platforms)]
             print(
-                f"{line}\n"
+                f"{OUTPUT_LINE}\n"
                 f"Note: Disabling multi-platform build, "
                 "this will only build a single-platform image.\n"
-                f"image: {sanitized_name} platform:{platforms[0]}\n"
-                f"{line}"
+                f"image: {repo} platform:{platforms[0]}\n"
+                f"{OUTPUT_LINE}"
             )
         else:
             print(
-                f"{line}\n"
+                f"{OUTPUT_LINE}\n"
                 f"Note: Building multi-platform images can take a long time, please be patient.\n"
                 "If you are running this locally, you can speed this up by using the "
                 "'--disable-multi-platform' CLI flag "
                 "or set the 'disable-multi-platform' flag in the global config file.\n"
-                f"{line}"
+                f"{OUTPUT_LINE}"
             )
 
         processes = []
         print(
             f'Starting builds for {len(platforms)} platforms in {"parallel" if do_multiprocessing else "sequence"}'
         )
+
+        # Since multiprocessing may be used, use a simple queue to communicate with the build method
+        # This queue receives tuples of (image_ref, image_digest)
+        image_info_by_image_ref = {}
+        queue = SimpleQueue()
         for platform in platforms:
-            platform_image_name = f"{base_image_name}-{platform.replace('/', '-')}"
+            tag = f"{built_image.id}-{platform.replace('/', '-')}"
+            image_ref = f"{repo}:{tag}"
+            # Contains all fields needed for the BuiltTaggedImage instance except digest, which will be added later
+            image_info_by_image_ref[image_ref] = {
+                "repo": repo,
+                "tag": tag,
+                "platform": platform,
+            }
             build_single_image_args = (
-                platform_image_name,
+                queue,
+                image_ref,
                 platform,
                 path,
                 dockerfile,
-                tags,
                 build_args,
-                mp_image_name,
                 inject,
                 cache,
                 pull,
             )
-            LOGGER.debug(f"Building {platform_image_name} for {platform}")
+            LOGGER.debug(f"Building {repo} for {platform}")
             if do_multiprocessing:
                 processes.append(
                     Process(
@@ -588,16 +458,31 @@ def get_path(file):
         for proc in processes:
             proc.join()
 
+        while not queue.empty():
+            image_ref, image_digest = queue.get()
+            assert (
+                image_ref in image_info_by_image_ref
+            ), f"Image ref {image_ref} is missing in generated info"
+            image_info = image_info_by_image_ref.pop(image_ref)
+            built_image.add_platform_image(
+                image_info.get("platform"),
+                BuiltTaggedImage(
+                    **image_info,
+                    digest=image_digest,
+                ),
+            )
+        assert not image_info_by_image_ref, f"Image refs were not generated successfully, unclaimed refs: {image_info_by_image_ref}"
+
         if cleanup_dockerfile and dockerfile and os.path.exists(dockerfile):
             os.remove(dockerfile)
 
-        return self._intermediate_built_images[mp_image_name]
+        return built_image
 
     @timeout_decorator.timeout(PUSH_TIMEOUT)
     def _push_with_timeout(self, src_names: List[str], tag_names: List[str]) -> None:
         """
         Creates tags from a set of source images in the remote registry.
-        This method will timeout if it takes too long. An exception may be
+        This method will time out if it takes too long. An exception may be
         caught and retried for the timeout.
 
         Args:
@@ -610,164 +495,93 @@ def _push_with_timeout(self, src_names: List[str], tag_names: List[str]) -> None
         LOGGER.info(f"Pushing sources {src_names} to tags {tag_names}")
         docker.buildx.imagetools.create(sources=src_names, tags=tag_names)
 
-    def push(self, name: str, dest_names: List[str] = None) -> List[str]:
+    def push(self, log: ConsoleLogger) -> 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
+        Pushes all built images to their tagged image counterparts.
+        :arg log: The log instance for output to the console.
+        :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
         retries = 5
 
-        src_images = self._intermediate_built_images[name]
-        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
-            # only need get tags for one image, since they should be identical
-            for tag in src_images[0].tags:
-                tagged_names.append(f"{dest_names}:{tag}")
-        else:
-            tagged_names = dest_names
+        for built_image in self._built_images:
+            if not built_image.tagged_images:
+                log.write(
+                    f"No tags exist for built image {built_image.id}, not pushing\n"
+                )
+                continue
 
-        for image in src_images:
-            for tag in image.tags:
-                src_names.append(f"{image.repo}:{tag}")
+            source_image_refs = [image.image_ref for image in built_image.built_images]
 
-        timeout_seconds = initial_timeout_seconds
-        while retries > 0:
-            retries -= 1
-            LOGGER.debug(
-                f"Creating manifest list {name} with timeout {timeout_seconds} seconds"
-            )
-            try:
-                # Push each tag individually in order to prevent strange errors with multiple matching tags
-                for tag_name in tagged_names:
-                    self._push_with_timeout(src_names, [tag_name])
-                # Process finished within timeout
-                LOGGER.info(f"Successfully created multiplatform images {dest_names}")
-                break
-            except Exception as exc:  # pylint: disable=broad-exception-caught
-                LOGGER.warning(
-                    f"Caught exception while pushing images, retrying: {exc}"
+            for tagged_image in built_image.tagged_images:
+                log.write(
+                    f"Pushing {built_image.id} to {', '.join(tagged_image.image_refs)}\n"
                 )
-            if retries == 0:
-                raise TimeoutError(
-                    f"Timeout pushing {dest_names} after {retries} retries"
-                    f" and {timeout_seconds} seconds each try"
-                )
-            timeout_seconds += timeout_step_seconds
 
-            # Cap timeout at max timeout
-            timeout_seconds = min(timeout_seconds, timeout_max_seconds)
-        return tagged_names
+                timeout_seconds = initial_timeout_seconds
+                while retries > 0:
+                    retries -= 1
+                    LOGGER.debug(
+                        f"Creating manifest(s) {tagged_image} with timeout {timeout_seconds} seconds"
+                    )
+                    try:
+                        # Push each tag individually in order to prevent strange errors with multiple matching tags
+                        for image_ref in tagged_image.image_refs:
+                            self._push_with_timeout(source_image_refs, [image_ref])
+                        # Process finished within timeout
+                        LOGGER.info(
+                            f"Successfully pushed multiplatform image(s) {tagged_image}"
+                        )
+                        break
+                    except Exception as exc:  # pylint: disable=broad-exception-caught
+                        LOGGER.warning(
+                            f"Caught exception while pushing images, retrying: {exc}"
+                        )
+                    if retries == 0:
+                        raise TimeoutError(
+                            f"Timeout pushing {tagged_image} after {retries} retries"
+                            f" and {timeout_seconds} seconds each try"
+                        )
+                    timeout_seconds += timeout_step_seconds
+
+                    # Cap timeout at max timeout
+                    timeout_seconds = min(timeout_seconds, timeout_max_seconds)
 
-    def _find_native_platform_images(self, name: str) -> Optional[ImageInfo]:
+    @property
+    def num_built_images(self) -> int:
         """
-        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
+        Returns the count of built images that currently exist.
         """
-        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
-            return None
-
-        match_platform = [
-            image
-            for image in self._intermediate_built_images[name]
-            if image.repo.endswith(pattern)
-        ]
-
-        # No matches found, change os
-        if not match_platform:
-            if host_os == "Darwin":
-                pattern = f"linux-{host_arch}"
-            match_platform = [
-                image
-                for image in self._intermediate_built_images[name]
-                if image.repo.endswith(pattern)
-            ]
-
-        assert (
-            len(match_platform) <= 1
-        ), f"Found more than one match for {name} and {pattern}: {match_platform}"
+        return len(self._built_images)
 
-        # Still no matches found, get the first image
-        if not match_platform:
-            return self._intermediate_built_images[name][0]
-
-        return match_platform[0]
-
-    def tag_single_platform(
-        self, name: str, tags: List[str] = None, dest_name: str = None
-    ) -> None:
+    @staticmethod
+    def tag_native_platform(built_image: BuiltImageInfo) -> None:
         """
-        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.
+        Tags a single built image into the host registry for the platform matching this machine's platform.
+        If a matching image to the native platform is not found, then the first image is tagged. The tags come
+        from all registered tagged images in the built image info.
+        :arg built_image: The built image to pull and tag
         """
-        # 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_images(name)
-        if dest_name is None:
-            dest_name = name
-
-        if self._tagged_images_names.get(name) is None:
-            self._tagged_images_names[name] = []
-
-        # Tags all the source images with the dest_name and tags
-        # Then removes the intermediate source images
-        for image in source_image.formatted_list():
+        for tagged_image in built_image.tagged_images:
+            LOGGER.debug(
+                f"Tagging {built_image} as {tagged_image} for the native platform"
+            )
+            native_image = built_image.native_platform_image
+
+            # Pull the native image and then tag it
             try:
-                docker.pull(image)
-                for tag in tags:
-                    dest_tag = f"{dest_name}:{tag}"
-                    docker.tag(image, dest_tag)
-                    LOGGER.debug(f"Tagged {image} as {dest_tag}")
-                    self._tagged_images_names[name].append(dest_tag)
-                docker.image.remove(image, force=True)
+                docker.pull(native_image.image_ref)
+                for tagged_image_ref in tagged_image.image_refs:
+                    docker.tag(native_image.image_ref, tagged_image_ref)
+                    LOGGER.debug(
+                        f"Tagged {native_image.image_ref} as {tagged_image_ref}"
+                    )
+                docker.image.remove(native_image.image_ref, force=True)
             except python_on_whales.exceptions.DockerException as err:
-                LOGGER.error(f"Failed while tagging {dest_name}: {err}")
+                LOGGER.error(
+                    f"Failed while tagging {native_image.image_ref} as {tagged_image}: {err}"
+                )
                 raise err
-
-    def get_built_images(self, name: str) -> List[str]:
-        """
-        Returns the list of tagged images for the given name
-
-        Args:
-            name (str): The name of the image to find
-
-        Returns:
-            List[str]: The list of built images for the given name
-        """
-        return self._intermediate_built_images[name]
diff --git a/buildrunner/steprunner/tasks/__init__.py b/buildrunner/steprunner/tasks/__init__.py
index 72cb5bf4..3b90ecd4 100644
--- a/buildrunner/steprunner/tasks/__init__.py
+++ b/buildrunner/steprunner/tasks/__init__.py
@@ -6,24 +6,6 @@
 with the terms of the Adobe license agreement accompanying it.
 """
 
-import re
-
-
-def sanitize_tag(tag, log=None):
-    """
-    Sanitize a tag to remove illegal characters.
-
-    :param tag: The tag to sanitize.
-    :param log: Optional log to write warnings to.
-    :return: The sanitized tag.
-    """
-    _tag = re.sub(r"[^-_\w.]+", "-", tag.lower())
-    if _tag != tag and log:
-        log.write(
-            f"Forcing tag to lowercase and removing illegal characters: {tag} => {_tag}\n"
-        )
-    return _tag
-
 
 class BuildStepRunnerTask:
     """
@@ -49,15 +31,3 @@ def cleanup(self, context):
         Subclasses override this method to perform any cleanup tasks.
         """
         pass
-
-
-class MultiPlatformBuildStepRunnerTask(BuildStepRunnerTask):
-    """
-    Base task class for tasks that need to build/use images for multiple platforms.
-    """
-
-    def get_unique_build_name(self):
-        """
-        Returns a unique build name for this build and step.
-        """
-        return f"{self.step_runner.name}-{sanitize_tag(self.step_runner.build_runner.build_id)}"
diff --git a/buildrunner/steprunner/tasks/build.py b/buildrunner/steprunner/tasks/build.py
index 5612ec3e..68b74700 100644
--- a/buildrunner/steprunner/tasks/build.py
+++ b/buildrunner/steprunner/tasks/build.py
@@ -17,11 +17,11 @@
     BuildRunnerConfigurationError,
     BuildRunnerProcessingError,
 )
-from buildrunner.steprunner.tasks import MultiPlatformBuildStepRunnerTask
+from buildrunner.steprunner.tasks import BuildStepRunnerTask
 from buildrunner.utils import is_dict
 
 
-class BuildBuildStepRunnerTask(MultiPlatformBuildStepRunnerTask):  # pylint: disable=too-many-instance-attributes
+class BuildBuildStepRunnerTask(BuildStepRunnerTask):  # pylint: disable=too-many-instance-attributes
     """
     Class used to manage "build" build tasks.
     """
@@ -245,28 +245,28 @@ def run(self, context):
         )
         try:
             if self.platforms:
-                built_platform_images = (
-                    self.step_runner.multi_platform.build_multiple_images(
-                        platforms=self.platforms,
-                        path=self.path,
-                        file=self.dockerfile,
-                        mp_image_name=self.get_unique_build_name(),
-                        build_args=self.buildargs,
-                        inject=self.to_inject,
-                        cache=not self.nocache,
-                        pull=self.pull,
-                    )
+                built_image = self.step_runner.multi_platform.build_multiple_images(
+                    platforms=self.platforms,
+                    path=self.path,
+                    file=self.dockerfile,
+                    build_args=self.buildargs,
+                    inject=self.to_inject,
+                    cache=not self.nocache,
+                    pull=self.pull,
                 )
 
-                number_of_images = len(self.platforms)
+                num_platforms = len(self.platforms)
 
                 if self.step_runner.multi_platform.disable_multi_platform:
-                    number_of_images = 1
+                    num_platforms = 1
+
+                num_built_platforms = len(built_image.platforms)
 
-                assert len(built_platform_images) == number_of_images, (
-                    f"Number of built images ({len(built_platform_images)}) does not match "
-                    f"the number of platforms ({number_of_images})"
+                assert num_built_platforms == num_platforms, (
+                    f"Number of built images ({num_built_platforms}) does not match "
+                    f"the number of platforms ({num_platforms})"
                 )
+                context["mp_built_image"] = built_image
 
             else:
                 exit_code = builder.build(
diff --git a/buildrunner/steprunner/tasks/push.py b/buildrunner/steprunner/tasks/push.py
index efd5c632..c9e3a909 100644
--- a/buildrunner/steprunner/tasks/push.py
+++ b/buildrunner/steprunner/tasks/push.py
@@ -6,6 +6,7 @@
 with the terms of the Adobe license agreement accompanying it.
 """
 import os
+import re
 from typing import List, Optional, Union
 
 import buildrunner.docker
@@ -13,7 +14,23 @@
     BuildRunnerConfigurationError,
     BuildRunnerProcessingError,
 )
-from buildrunner.steprunner.tasks import MultiPlatformBuildStepRunnerTask, sanitize_tag
+from buildrunner.steprunner.tasks import BuildStepRunnerTask
+
+
+def sanitize_tag(tag, log=None):
+    """
+    Sanitize a tag to remove illegal characters.
+
+    :param tag: The tag to sanitize.
+    :param log: Optional log to write warnings to.
+    :return: The sanitized tag.
+    """
+    _tag = re.sub(r"[^-_\w.]+", "-", tag.lower())
+    if _tag != tag and log:
+        log.write(
+            f"Forcing tag to lowercase and removing illegal characters: {tag} => {_tag}\n"
+        )
+    return _tag
 
 
 class RepoDefinition:
@@ -51,7 +68,7 @@ def __init__(
             self.repository = self.repository[0:tag_index]
 
 
-class PushBuildStepRunnerTask(MultiPlatformBuildStepRunnerTask):
+class PushBuildStepRunnerTask(BuildStepRunnerTask):
     """
     Class used to push the resulting image (either from the build task, or if
     there is a run task, the snapshot of the resulting run container) to the
@@ -89,43 +106,42 @@ def run(self, context):  # pylint: disable=too-many-branches
         default_tag = sanitize_tag(self.step_runner.build_runner.build_id)
 
         # Tag multi-platform images
-        if self.step_runner.multi_platform.is_multiplatform(
-            self.get_unique_build_name()
-        ):
+        built_image = context.get("mp_built_image")
+        if built_image:
+            # These are used in the image artifacts below, and should match for all tagged images
+            built_image_ids_str = ",".join(
+                [image.trunc_digest for image in built_image.built_images]
+            )
+            built_image_id_with_platforms = [
+                f"{image.platform}:{image.trunc_digest}"
+                for image in built_image.built_images
+            ]
+
             for repo in self._repos:
-                # Always add default tag
+                # Always add default tag and then add it to the built image info
                 repo.tags.append(default_tag)
-                self.step_runner.multi_platform.tag_single_platform(
-                    name=self.get_unique_build_name(),
-                    tags=repo.tags,
-                    dest_name=repo.repository,
-                )
+                tagged_image = built_image.add_tagged_image(repo.repository, repo.tags)
 
-                for tag in repo.tags:
-                    self.step_runner.build_runner.committed_images.add(
-                        f"{repo.repository}:{tag}"
-                    )
+                # Add tagged image refs to committed images for use in determining if pull should be true/false
+                for image_ref in tagged_image.image_refs:
+                    self.step_runner.build_runner.committed_images.add(image_ref)
 
-                # add image as artifact
+                # Add tagged image as artifact if this is a push and not just a commit
                 if not self._commit_only:
-                    images = self.step_runner.multi_platform.get_built_images(
-                        self.get_unique_build_name()
-                    )
-                    image_ids = ",".join([image.trunc_digest() for image in images])
-                    platforms = [
-                        f"{image.platform}:{image.trunc_digest()}" for image in images
-                    ]
                     self.step_runner.build_runner.add_artifact(
                         repo.repository,
                         {
                             "type": "docker-image",
-                            "docker:image": image_ids,
+                            "docker:image": built_image_ids_str,
                             "docker:repository": repo.repository,
                             "docker:tags": repo.tags,
-                            "docker:platforms": platforms,
+                            "docker:platforms": built_image_id_with_platforms,
                         },
                     )
 
+            # Tag all images locally for the native platform
+            self.step_runner.multi_platform.tag_native_platform(built_image)
+
         # Tag single platform images
         else:
             # first see if a run task produced an image (via a post-build config)
diff --git a/tests/test_multiplatform.py b/tests/test_multiplatform.py
index 2066a45e..285fa11e 100644
--- a/tests/test_multiplatform.py
+++ b/tests/test_multiplatform.py
@@ -6,10 +6,9 @@
 from python_on_whales import docker
 from python_on_whales.exceptions import DockerException
 
-from buildrunner.docker.multiplatform_image_builder import (
-    ImageInfo,
-    MultiplatformImageBuilder,
-)
+from buildrunner.docker.multiplatform_image_builder import MultiplatformImageBuilder
+from buildrunner.docker.image_info import BuiltImageInfo
+
 
 TEST_DIR = os.path.basename(os.path.dirname(__file__))
 
@@ -30,26 +29,25 @@ def _get_uuid():
         yield uuid_mock
 
 
-def actual_images_match_expected(actual_images, expected_images) -> List[str]:
+def _actual_images_match_expected(
+    built_image: BuiltImageInfo, expected_tags
+) -> List[str]:
+    actual_images = built_image.built_images
     missing_images = []
-    found = False
-    for expected_image in expected_images:
+    for expected_tag in expected_tags:
         found = False
         for actual_image in actual_images:
-            if actual_image.repo.endswith(expected_image):
+            if actual_image.tag.endswith(expected_tag):
                 found = True
         if not found:
-            missing_images.append(expected_image)
+            missing_images.append(expected_tag)
     return missing_images
 
 
 def test_start_local_registry():
-    registry_name = None
-    volume_name = None
-
-    with MultiplatformImageBuilder() as mp:
-        mp._start_local_registry()
-        registry_name = mp._mp_registry_info.name
+    with MultiplatformImageBuilder() as mpib:
+        mpib._start_local_registry()
+        registry_name = mpib._mp_registry_info.name
 
         # Check that the registry is running and only one is found with that name
         registry_container = docker.ps(filters={"name": registry_name})
@@ -75,18 +73,14 @@ def test_start_local_registry():
 
 
 def test_start_local_registry_on_build():
-    registry_name = None
-    volume_name = None
-
-    with MultiplatformImageBuilder() as mp:
+    with MultiplatformImageBuilder() as mpib:
         # Check that the registry is NOT running
-        assert mp._mp_registry_info is None
-        assert mp._local_registry_is_running is False
+        assert mpib._mp_registry_info is None
+        assert mpib._local_registry_is_running is False
 
         # Building should start the registry
         test_path = f"{TEST_DIR}/test-files/multiplatform"
-        mp.build_multiple_images(
-            mp_image_name="test-images-2000-start-on-build",
+        mpib.build_multiple_images(
             platforms=["linux/arm64", "linux/amd64"],
             path=test_path,
             file=f"{test_path}/Dockerfile",
@@ -94,7 +88,7 @@ def test_start_local_registry_on_build():
         )
 
         # Check that the registry is running and only one is found with that name
-        registry_name = mp._mp_registry_info.name
+        registry_name = mpib._mp_registry_info.name
         first_registry_name = registry_name
         registry_container = docker.ps(filters={"name": registry_name})
         assert len(registry_container) == 1
@@ -113,15 +107,14 @@ def test_start_local_registry_on_build():
         assert registry_container.state.running
 
         # Building again should not start a new registry
-        mp.build_multiple_images(
-            mp_image_name="test-images-2000-start-on-build2",
+        mpib.build_multiple_images(
             platforms=["linux/arm64", "linux/amd64"],
             path=test_path,
             file=f"{test_path}/Dockerfile",
             do_multiprocessing=False,
         )
 
-        registry_name = mp._mp_registry_info.name
+        registry_name = mpib._mp_registry_info.name
         assert first_registry_name == registry_name
 
     # Check that the registry is stopped and cleaned up
@@ -130,145 +123,40 @@ def test_start_local_registry_on_build():
     assert not docker.volume.exists(volume_name)
 
 
+# TODO Fix tests
 @pytest.mark.parametrize(
-    "name, in_mock_os, in_mock_arch, built_images, expected_image",
-    [
-        # platform = linux/arm64
-        (
-            "test-images-2000",
-            "linux",
-            "arm64",
-            {
-                "test-images-2000": [
-                    ImageInfo(
-                        "localhost:32828/test-images-2000-linux-amd64", ["latest"]
-                    ),
-                    ImageInfo(
-                        "localhost:32828/test-images-2000-linux-arm64", ["latest"]
-                    ),
-                ]
-            },
-            "localhost:32828/test-images-2000-linux-arm64:latest",
-        ),
-        # OS does not match for Darwin change to linux
-        (
-            "test-images-2000",
-            "Darwin",
-            "arm64",
-            {
-                "test-images-2000": [
-                    ImageInfo(
-                        "localhost:32829/test-images-2000-linux-amd64", ["latest"]
-                    ),
-                    ImageInfo(
-                        "localhost:32829/test-images-2000-linux-arm64", ["latest"]
-                    ),
-                ]
-            },
-            "localhost:32829/test-images-2000-linux-arm64:latest",
-        ),
-        # platform = linux/amd64
-        (
-            "test-images-2000",
-            "linux",
-            "amd64",
-            {
-                "test-images-2000": [
-                    ImageInfo(
-                        "localhost:32811/test-images-2000-linux-amd64", ["latest"]
-                    ),
-                    ImageInfo(
-                        "localhost:32811/test-images-2000-linux-arm64", ["latest"]
-                    ),
-                ]
-            },
-            "localhost:32811/test-images-2000-linux-amd64:latest",
-        ),
-        # No match found, get the first image
-        (
-            "test-images-2000",
-            "linux",
-            "arm",
-            {
-                "test-images-2000": [
-                    ImageInfo(
-                        "localhost:32830/test-images-2000-linux-amd64", ["0.1.0"]
-                    ),
-                    ImageInfo(
-                        "localhost:32830/test-images-2000-linux-amd64", ["0.2.0"]
-                    ),
-                ]
-            },
-            "localhost:32830/test-images-2000-linux-amd64:0.1.0",
-        ),
-        # Built_images for name does not exist in dictionary
-        (
-            "test-images-2001",
-            "linux",
-            "arm64",
-            {
-                "test-images-2000": [
-                    ImageInfo(
-                        "localhost:32831/test-images-2000-linux-amd64", ["latest"]
-                    ),
-                    ImageInfo(
-                        "localhost:32831/test-images-2000-linux-arm64", ["latest"]
-                    ),
-                ]
-            },
-            None,
-        ),
-        # Built_images for name is empty
-        ("test-images-2000", "linux", "arm64", {"test-images-2000": []}, None),
-    ],
-)
-@patch("buildrunner.docker.multiplatform_image_builder.machine")
-@patch("buildrunner.docker.multiplatform_image_builder.system")
-def test_find_native_platform(
-    mock_os, mock_arch, name, in_mock_os, in_mock_arch, built_images, expected_image
-):
-    mock_os.return_value = in_mock_os
-    mock_arch.return_value = in_mock_arch
-    with MultiplatformImageBuilder() as mp:
-        mp._intermediate_built_images = built_images
-        found_platform = mp._find_native_platform_images(name)
-        assert str(found_platform) == str(expected_image)
-
-
-@pytest.mark.parametrize(
-    "name, platforms, expected_image_names",
-    [("test-image-tag-2000", ["linux/arm64"], ["buildrunner-mp-uuid1-linux-arm64"])],
+    "name, platforms, expected_image_tags",
+    [("test-image-tag-2000", ["linux/arm64"], ["uuid1-linux-arm64"])],
 )
-def test_tag_single_platform(name, platforms, expected_image_names):
+def test_tag_native_platform(name, platforms, expected_image_tags):
     tag = "latest"
     test_path = f"{TEST_DIR}/test-files/multiplatform"
-    with MultiplatformImageBuilder(cleanup_images=True) as mp:
-        built_images = mp.build_multiple_images(
-            mp_image_name=name,
-            platforms=platforms,
+    with MultiplatformImageBuilder(cleanup_images=True) as mpib:
+        built_image = mpib.build_multiple_images(
+            platforms,
             path=test_path,
             file=f"{test_path}/Dockerfile",
             do_multiprocessing=False,
         )
 
         assert (
-            built_images is not None and len(built_images) == 1
+            built_image is not None and len(built_image.built_images) == 1
         ), f"Failed to build {name} for {platforms}"
-        missing_images = actual_images_match_expected(
-            built_images, expected_image_names
-        )
+        missing_images = _actual_images_match_expected(built_image, expected_image_tags)
         assert (
             missing_images == []
-        ), f"Failed to find {missing_images} in {[image.repo for image in built_images]}"
+        ), f"Failed to find {missing_images} in {[image.repo for image in built_image.built_images]}"
 
-        mp.tag_single_platform(name)
+        mpib.tag_native_platform(built_image)
         # Check that the image was tagged and present
         found_image = docker.image.list(filters={"reference": f"{name}:{tag}"})
         assert len(found_image) == 1
         assert f"{name}:{tag}" in found_image[0].repo_tags
 
         # Check that intermediate images are not on host registry
-        found_image = docker.image.list(filters={"reference": f"{built_images[0]}*"})
+        found_image = docker.image.list(
+            filters={"reference": f"{built_image.built_images[0].image_ref}*"}
+        )
         assert len(found_image) == 0
 
     # Check that the image has be removed for host registry
@@ -277,15 +165,14 @@ def test_tag_single_platform(name, platforms, expected_image_names):
 
 
 @pytest.mark.parametrize(
-    "name, platforms, expected_image_names",
-    [("test-image-tag-2000", ["linux/arm64"], ["buildrunner-mp-uuid1-linux-arm64"])],
+    "name, platforms, expected_image_tags",
+    [("test-image-tag-2000", ["linux/arm64"], ["uuid1-linux-arm64"])],
 )
-def test_tag_single_platform_multiple_tags(name, platforms, expected_image_names):
+def test_tag_native_platform_multiple_tags(name, platforms, expected_image_tags):
     tags = ["latest", "0.1.0"]
     test_path = f"{TEST_DIR}/test-files/multiplatform"
-    with MultiplatformImageBuilder(cleanup_images=True) as mp:
-        built_images = mp.build_multiple_images(
-            mp_image_name=name,
+    with MultiplatformImageBuilder(cleanup_images=True) as mpib:
+        built_image = mpib.build_multiple_images(
             platforms=platforms,
             path=test_path,
             file=f"{test_path}/Dockerfile",
@@ -293,16 +180,15 @@ def test_tag_single_platform_multiple_tags(name, platforms, expected_image_names
         )
 
         assert (
-            built_images is not None and len(built_images) == 1
+            built_image is not None and len(built_image.built_images) == 1
         ), f"Failed to build {name} for {platforms}"
-        missing_images = actual_images_match_expected(
-            built_images, expected_image_names
-        )
+        missing_images = _actual_images_match_expected(built_image, expected_image_tags)
         assert (
             missing_images == []
-        ), f"Failed to find {missing_images} in {[image.repo for image in built_images]}"
+        ), f"Failed to find {missing_images} in {[image.repo for image in built_image.built_images]}"
 
-        mp.tag_single_platform(name=name, tags=tags)
+        built_image.add_tagged_image(repo=name, tags=tags)
+        mpib.tag_native_platform(built_image)
         # Check that the image was tagged and present
         found_image = docker.image.list(filters={"reference": f"{name}*"})
         assert len(found_image) == 1
@@ -310,7 +196,9 @@ def test_tag_single_platform_multiple_tags(name, platforms, expected_image_names
             assert f"{name}:{tag}" in found_image[0].repo_tags
 
         # Check that intermediate images are not on host registry
-        found_image = docker.image.list(filters={"reference": f"{built_images[0]}*"})
+        found_image = docker.image.list(
+            filters={"reference": f"{built_image.built_images[0].image_ref}*"}
+        )
         assert len(found_image) == 0
 
     # Check that the tagged image has be removed for host registry
@@ -319,16 +207,15 @@ def test_tag_single_platform_multiple_tags(name, platforms, expected_image_names
 
 
 @pytest.mark.parametrize(
-    "name, platforms, expected_image_names",
-    [("test-image-tag-2000", ["linux/arm64"], ["buildrunner-mp-uuid1-linux-arm64"])],
+    "name, platforms, expected_image_tags",
+    [("test-image-tag-2000", ["linux/arm64"], ["uuid1-linux-arm64"])],
 )
-def test_tag_single_platform_keep_images(name, platforms, expected_image_names):
+def test_tag_native_platform_keep_images(name, platforms, expected_image_tags):
     tag = "latest"
     test_path = f"{TEST_DIR}/test-files/multiplatform"
     try:
-        with MultiplatformImageBuilder(cleanup_images=False) as mp:
-            built_images = mp.build_multiple_images(
-                mp_image_name=name,
+        with MultiplatformImageBuilder(cleanup_images=False) as mpib:
+            built_image = mpib.build_multiple_images(
                 platforms=platforms,
                 path=test_path,
                 file=f"{test_path}/Dockerfile",
@@ -336,16 +223,17 @@ def test_tag_single_platform_keep_images(name, platforms, expected_image_names):
             )
 
             assert (
-                built_images is not None and len(built_images) == 1
+                built_image is not None and len(built_image.built_images) == 1
             ), f"Failed to build {name} for {platforms}"
-            missing_images = actual_images_match_expected(
-                built_images, expected_image_names
+            missing_images = _actual_images_match_expected(
+                built_image, expected_image_tags
             )
             assert (
                 missing_images == []
-            ), f"Failed to find {missing_images} in {[image.repo for image in built_images]}"
+            ), f"Failed to find {missing_images} in {[image.repo for image in built_image.built_images]}"
 
-            mp.tag_single_platform(name)
+            built_image.add_tagged_image(repo=name, tags=["latest"])
+            mpib.tag_native_platform(built_image)
 
             # Check that the image was tagged and present
             found_image = docker.image.list(filters={"reference": f"{name}:{tag}"})
@@ -354,7 +242,7 @@ def test_tag_single_platform_keep_images(name, platforms, expected_image_names):
 
             # Check that intermediate images are not on host registry
             found_image = docker.image.list(
-                filters={"reference": f"{built_images[0]}*"}
+                filters={"reference": f"{built_image.built_images[0].image_ref}*"}
             )
             assert len(found_image) == 0
 
@@ -377,18 +265,17 @@ def test_push():
             platforms = ["linux/arm64", "linux/amd64"]
 
             test_path = f"{TEST_DIR}/test-files/multiplatform"
-            with MultiplatformImageBuilder() as mp:
-                built_images = mp.build_multiple_images(
-                    mp_image_name=build_name,
+            with MultiplatformImageBuilder() as mpib:
+                built_image = mpib.build_multiple_images(
                     platforms=platforms,
                     path=test_path,
                     file=f"{test_path}/Dockerfile",
                     do_multiprocessing=False,
-                    tags=tags,
                 )
 
-                assert built_images is not None
-                mp.push(build_name)
+                assert built_image is not None
+                built_image.add_tagged_image(repo=build_name, tags=tags)
+                mpib.push(MagicMock())
 
                 # Make sure the image isn't in the local registry
                 docker.image.remove(build_name, force=True)
@@ -410,7 +297,6 @@ def test_push():
 
 def test_push_with_dest_names():
     dest_names = None
-    built_images = None
     try:
         with MultiplatformImageBuilder() as remote_mp:
             remote_mp._start_local_registry()
@@ -423,18 +309,18 @@ def test_push_with_dest_names():
             platforms = ["linux/arm64", "linux/amd64"]
 
             test_path = f"{TEST_DIR}/test-files/multiplatform"
-            with MultiplatformImageBuilder() as mp:
-                built_images = mp.build_multiple_images(
-                    mp_image_name=build_name,
+            with MultiplatformImageBuilder() as mpib:
+                built_image = mpib.build_multiple_images(
                     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)
+                assert built_image is not None
+                for dest_name in dest_names:
+                    built_image.add_tagged_image(repo=dest_name, tags=tags)
+                mpib.push(MagicMock())
 
                 # Make sure the image isn't in the local registry
                 for dest_name in dest_names:
@@ -457,17 +343,17 @@ def test_push_with_dest_names():
 
 
 @pytest.mark.parametrize(
-    "name, platforms, expected_image_names",
+    "name, platforms, expected_image_tags",
     [
         (
             "test-build-image-2000",
             ["linux/arm64"],
-            ["buildrunner-mp-uuid1-linux-arm64"],
+            ["uuid1-linux-arm64"],
         ),
         (
             "test-build-image-2001",
             ["linux/amd64", "linux/arm64"],
-            ["buildrunner-mp-uuid1-linux-amd64", "buildrunner-mp-uuid1-linux-arm64"],
+            ["uuid1-linux-amd64", "uuid1-linux-arm64"],
         ),
     ],
 )
@@ -484,29 +370,30 @@ def test_build(
     mock_remove,
     name,
     platforms,
-    expected_image_names,
+    expected_image_tags,
 ):
+    _ = mock_build
+    _ = mock_pull
+    _ = mock_push
+    _ = mock_remove
     mock_inspect.return_value = MagicMock()
     mock_inspect.return_value.id = "myfakeimageid"
     test_path = f"{TEST_DIR}/test-files/multiplatform"
-    with MultiplatformImageBuilder() as mp:
-        built_images = mp.build_multiple_images(
-            mp_image_name=name,
+    with MultiplatformImageBuilder() as mpib:
+        built_image = mpib.build_multiple_images(
             platforms=platforms,
             path=test_path,
             file=f"{test_path}/Dockerfile",
             do_multiprocessing=False,
         )
 
-        assert len(built_images) == len(platforms)
-        assert len(built_images) == len(expected_image_names)
+        assert len(built_image.built_images) == len(platforms)
+        assert len(built_image.built_images) == len(expected_image_tags)
 
-        missing_images = actual_images_match_expected(
-            built_images, expected_image_names
-        )
+        missing_images = _actual_images_match_expected(built_image, expected_image_tags)
         assert (
             missing_images == []
-        ), f"Failed to find {missing_images} in {[image.repo for image in built_images]}"
+        ), f"Failed to find {missing_images} in {[image.repo for image in built_image.built_images]}"
 
 
 @patch("buildrunner.docker.multiplatform_image_builder.docker.image.remove")
@@ -517,27 +404,25 @@ def test_build(
 def test_build_multiple_builds(
     mock_build, mock_pull, mock_push, mock_inspect, mock_remove
 ):
+    _ = mock_remove
     mock_inspect.return_value = MagicMock()
     mock_inspect.return_value.id = "myfakeimageid"
-    name1 = "test-build-multi-image-2001"
     platforms1 = ["linux/amd64", "linux/arm64"]
-    expected_image_names1 = [
-        "buildrunner-mp-uuid1-linux-amd64",
-        "buildrunner-mp-uuid1-linux-arm64",
+    expected_image_tags1 = [
+        "uuid1-linux-amd64",
+        "uuid1-linux-arm64",
     ]
 
-    name2 = "test-build-multi-image-2002"
     platforms2 = ["linux/amd64", "linux/arm64"]
-    expected_image_names2 = [
-        "buildrunner-mp-uuid2-linux-amd64",
-        "buildrunner-mp-uuid2-linux-arm64",
+    expected_image_tags2 = [
+        "uuid2-linux-amd64",
+        "uuid2-linux-arm64",
     ]
 
     test_path = f"{TEST_DIR}/test-files/multiplatform"
-    with MultiplatformImageBuilder() as mp:
+    with MultiplatformImageBuilder() as mpib:
         # Build set 1
-        built_images1 = mp.build_multiple_images(
-            mp_image_name=name1,
+        built_image1 = mpib.build_multiple_images(
             platforms=platforms1,
             path=test_path,
             file=f"{test_path}/Dockerfile",
@@ -545,8 +430,7 @@ def test_build_multiple_builds(
         )
 
         # Build set 2
-        built_images2 = mp.build_multiple_images(
-            mp_image_name=name2,
+        built_image2 = mpib.build_multiple_images(
             platforms=platforms2,
             path=test_path,
             file=f"{test_path}/Dockerfile",
@@ -554,31 +438,31 @@ def test_build_multiple_builds(
         )
 
         # Check set 1
-        assert len(built_images1) == len(platforms1)
-        assert len(built_images1) == len(expected_image_names1)
-        missing_images = actual_images_match_expected(
-            built_images1, expected_image_names1
+        assert len(built_image1.built_images) == len(platforms1)
+        assert len(built_image1.built_images) == len(expected_image_tags1)
+        missing_images = _actual_images_match_expected(
+            built_image1, expected_image_tags1
         )
         assert (
             missing_images == []
-        ), f"Failed to find {missing_images} in {[image.repo for image in built_images1]}"
+        ), f"Failed to find {missing_images} in {[image.repo for image in built_image1.built_images1]}"
 
         # Check set 2
-        assert len(built_images2) == len(platforms2)
-        assert len(built_images2) == len(expected_image_names2)
-        missing_images = actual_images_match_expected(
-            built_images2, expected_image_names2
+        assert len(built_image2.built_images) == len(platforms2)
+        assert len(built_image2.built_images) == len(expected_image_tags2)
+        missing_images = _actual_images_match_expected(
+            built_image2, expected_image_tags2
         )
         assert (
             missing_images == []
-        ), f"Failed to find {missing_images} in {[image.repo for image in built_images2]}"
+        ), f"Failed to find {missing_images} in {[image.repo for image in built_image2.built_images]}"
 
     assert mock_build.call_count == 4
     prefix = mock_build.call_args.kwargs["tags"][0].split("buildrunner-mp")[0]
     assert mock_build.call_args_list == [
         call(
             "tests/test-files/multiplatform",
-            tags=[f"{prefix}buildrunner-mp-uuid1-linux-amd64:latest"],
+            tags=[f"{prefix}buildrunner-mp:uuid1-linux-amd64"],
             platforms=["linux/amd64"],
             load=True,
             file="tests/test-files/multiplatform/Dockerfile",
@@ -591,7 +475,7 @@ def test_build_multiple_builds(
         ),
         call(
             "tests/test-files/multiplatform",
-            tags=[f"{prefix}buildrunner-mp-uuid1-linux-arm64:latest"],
+            tags=[f"{prefix}buildrunner-mp:uuid1-linux-arm64"],
             platforms=["linux/arm64"],
             load=True,
             file="tests/test-files/multiplatform/Dockerfile",
@@ -604,7 +488,7 @@ def test_build_multiple_builds(
         ),
         call(
             "tests/test-files/multiplatform",
-            tags=[f"{prefix}buildrunner-mp-uuid2-linux-amd64:latest"],
+            tags=[f"{prefix}buildrunner-mp:uuid2-linux-amd64"],
             platforms=["linux/amd64"],
             load=True,
             file="tests/test-files/multiplatform/Dockerfile",
@@ -617,7 +501,7 @@ def test_build_multiple_builds(
         ),
         call(
             "tests/test-files/multiplatform",
-            tags=[f"{prefix}buildrunner-mp-uuid2-linux-arm64:latest"],
+            tags=[f"{prefix}buildrunner-mp:uuid2-linux-arm64"],
             platforms=["linux/arm64"],
             load=True,
             file="tests/test-files/multiplatform/Dockerfile",
@@ -631,68 +515,20 @@ def test_build_multiple_builds(
     ]
     assert mock_push.call_count == 4
     assert mock_push.call_args_list == [
-        call([f"{prefix}buildrunner-mp-uuid1-linux-amd64:latest"]),
-        call([f"{prefix}buildrunner-mp-uuid1-linux-arm64:latest"]),
-        call([f"{prefix}buildrunner-mp-uuid2-linux-amd64:latest"]),
-        call([f"{prefix}buildrunner-mp-uuid2-linux-arm64:latest"]),
+        call([f"{prefix}buildrunner-mp:uuid1-linux-amd64"]),
+        call([f"{prefix}buildrunner-mp:uuid1-linux-arm64"]),
+        call([f"{prefix}buildrunner-mp:uuid2-linux-amd64"]),
+        call([f"{prefix}buildrunner-mp:uuid2-linux-arm64"]),
     ]
     assert mock_pull.call_count == 4
     assert mock_pull.call_args_list == [
-        call(f"{prefix}buildrunner-mp-uuid1-linux-amd64:latest"),
-        call(f"{prefix}buildrunner-mp-uuid1-linux-arm64:latest"),
-        call(f"{prefix}buildrunner-mp-uuid2-linux-amd64:latest"),
-        call(f"{prefix}buildrunner-mp-uuid2-linux-arm64:latest"),
+        call(f"{prefix}buildrunner-mp:uuid1-linux-amd64"),
+        call(f"{prefix}buildrunner-mp:uuid1-linux-arm64"),
+        call(f"{prefix}buildrunner-mp:uuid2-linux-amd64"),
+        call(f"{prefix}buildrunner-mp:uuid2-linux-arm64"),
     ]
 
 
-@pytest.mark.parametrize(
-    "name, tags, platforms, expected_image_names",
-    [
-        (
-            "test-build-tag-image-2000",
-            ["latest", "0.1.0"],
-            ["linux/arm64"],
-            ["buildrunner-mp-uuid1-linux-arm64"],
-        ),
-        (
-            "test-build-tag-image-2001",
-            ["latest", "0.2.0"],
-            ["linux/amd64", "linux/arm64"],
-            ["buildrunner-mp-uuid1-linux-amd64", "buildrunner-mp-uuid1-linux-arm64"],
-        ),
-    ],
-)
-def test_build_with_tags(name, tags, platforms, expected_image_names):
-    test_path = f"{TEST_DIR}/test-files/multiplatform"
-    with MultiplatformImageBuilder() as mp:
-        built_images = mp.build_multiple_images(
-            mp_image_name=name,
-            platforms=platforms,
-            path=test_path,
-            file=f"{test_path}/Dockerfile",
-            tags=tags,
-            do_multiprocessing=False,
-        )
-
-        assert len(built_images) == len(platforms)
-        assert len(built_images) == len(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]}"
-
-
-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
-
-
 @pytest.mark.parametrize(
     "builder, cache_builders, return_cache_options",
     [

From 68c5c629073336ebbba388891992c9e217fc3af7 Mon Sep 17 00:00:00 2001
From: saville <saville@adobe.com>
Date: Fri, 22 Dec 2023 08:45:53 -0700
Subject: [PATCH 2/3] Fix tests and add tests for image info native image
 retrieval

---
 buildrunner/__init__.py                       |  1 -
 .../docker/multiplatform_image_builder.py     |  2 -
 tests/test_image_info.py                      | 43 ++++++++++
 tests/test_multiplatform.py                   | 80 +++++++++----------
 4 files changed, 79 insertions(+), 47 deletions(-)
 create mode 100644 tests/test_image_info.py

diff --git a/buildrunner/__init__.py b/buildrunner/__init__.py
index 8be3f301..d426dbd5 100644
--- a/buildrunner/__init__.py
+++ b/buildrunner/__init__.py
@@ -606,7 +606,6 @@ def run(self):  # pylint: disable=too-many-statements,too-many-branches,too-many
         try:  # pylint: disable=too-many-nested-blocks
             with MultiplatformImageBuilder(
                 docker_registry=self.global_config.get_docker_registry(),
-                cleanup_images=self.cleanup_images,
                 temp_dir=self.global_config.get_temp_dir(),
                 disable_multi_platform=self.disable_multi_platform,
                 platform_builders=self.global_config.get("platform-builders"),
diff --git a/buildrunner/docker/multiplatform_image_builder.py b/buildrunner/docker/multiplatform_image_builder.py
index 7899939d..44a259a3 100644
--- a/buildrunner/docker/multiplatform_image_builder.py
+++ b/buildrunner/docker/multiplatform_image_builder.py
@@ -70,7 +70,6 @@ def __init__(
         self,
         docker_registry: Optional[str] = None,
         use_local_registry: bool = True,
-        cleanup_images: bool = False,
         temp_dir: str = os.getcwd(),
         disable_multi_platform: bool = False,
         platform_builders: Optional[Dict[str, str]] = None,
@@ -81,7 +80,6 @@ def __init__(
         self._docker_registry = docker_registry
         self._mp_registry_info = None
         self._use_local_registry = use_local_registry
-        self._cleanup_images = cleanup_images
         self._temp_dir = temp_dir
         self._disable_multi_platform = disable_multi_platform
         self._platform_builders = platform_builders
diff --git a/tests/test_image_info.py b/tests/test_image_info.py
new file mode 100644
index 00000000..2623d0a1
--- /dev/null
+++ b/tests/test_image_info.py
@@ -0,0 +1,43 @@
+from unittest import mock
+
+import pytest
+
+from buildrunner.docker.image_info import BuiltTaggedImage, BuiltImageInfo
+
+
+@pytest.mark.parametrize(
+    "system_result, machine_result, platform_result",
+    [
+        ("Darwin", "x86_64", "linux/amd64"),
+        ("Darwin", "aarch64", "linux/arm64/v7"),
+        ("Linux", "x86_64", "linux/amd64"),
+        ("Linux", "aarch64", "linux/arm64/v7"),
+        ("Bogus", "Noarch", "bogus"),
+    ],
+)
+@mock.patch("buildrunner.docker.image_info.python_platform")
+def test_find_native_platform(
+    platform_mock,
+    system_result,
+    machine_result,
+    platform_result,
+):
+    platform_mock.system.return_value = system_result
+    platform_mock.machine.return_value = machine_result
+
+    run_id = "abc123"
+    built_image = BuiltImageInfo(id=run_id)
+    for platform in ("bogus", "linux/arm64/v7", "linux/amd64"):
+        built_image.add_platform_image(
+            platform,
+            BuiltTaggedImage(
+                repo="repo1",
+                tag=f'{run_id}-{platform.replace("/", "-")}',
+                digest="12345",
+                platform=platform,
+            ),
+        )
+    assert (
+        built_image.native_platform_image
+        == built_image.images_by_platform[platform_result]
+    )
diff --git a/tests/test_multiplatform.py b/tests/test_multiplatform.py
index 285fa11e..8b04f849 100644
--- a/tests/test_multiplatform.py
+++ b/tests/test_multiplatform.py
@@ -10,7 +10,7 @@
 from buildrunner.docker.image_info import BuiltImageInfo
 
 
-TEST_DIR = os.path.basename(os.path.dirname(__file__))
+TEST_DIR = os.path.dirname(__file__)
 
 # FIXME: These tests can be broken if a custom buildx builder is set as default  # pylint: disable=fixme
 
@@ -123,15 +123,13 @@ def test_start_local_registry_on_build():
     assert not docker.volume.exists(volume_name)
 
 
-# TODO Fix tests
 @pytest.mark.parametrize(
-    "name, platforms, expected_image_tags",
-    [("test-image-tag-2000", ["linux/arm64"], ["uuid1-linux-arm64"])],
+    "platforms, expected_image_tags",
+    [(["linux/arm64"], ["uuid1-linux-arm64"])],
 )
-def test_tag_native_platform(name, platforms, expected_image_tags):
-    tag = "latest"
+def test_tag_native_platform(platforms, expected_image_tags):
     test_path = f"{TEST_DIR}/test-files/multiplatform"
-    with MultiplatformImageBuilder(cleanup_images=True) as mpib:
+    with MultiplatformImageBuilder() as mpib:
         built_image = mpib.build_multiple_images(
             platforms,
             path=test_path,
@@ -141,7 +139,7 @@ def test_tag_native_platform(name, platforms, expected_image_tags):
 
         assert (
             built_image is not None and len(built_image.built_images) == 1
-        ), f"Failed to build {name} for {platforms}"
+        ), f"Failed to build for {platforms}"
         missing_images = _actual_images_match_expected(built_image, expected_image_tags)
         assert (
             missing_images == []
@@ -149,19 +147,17 @@ def test_tag_native_platform(name, platforms, expected_image_tags):
 
         mpib.tag_native_platform(built_image)
         # Check that the image was tagged and present
-        found_image = docker.image.list(filters={"reference": f"{name}:{tag}"})
+        found_image = docker.image.list(
+            filters={"reference": built_image.built_images[0].image_ref}
+        )
         assert len(found_image) == 1
-        assert f"{name}:{tag}" in found_image[0].repo_tags
+        assert built_image.built_images[0].image_ref in found_image[0].repo_tags
 
-        # Check that intermediate images are not on host registry
+        # Check that intermediate images are on the host registry
         found_image = docker.image.list(
             filters={"reference": f"{built_image.built_images[0].image_ref}*"}
         )
-        assert len(found_image) == 0
-
-    # Check that the image has be removed for host registry
-    found_image = docker.image.list(filters={"reference": f"{name}:{tag}"})
-    assert len(found_image) == 0
+        assert len(found_image) == 1
 
 
 @pytest.mark.parametrize(
@@ -171,7 +167,7 @@ def test_tag_native_platform(name, platforms, expected_image_tags):
 def test_tag_native_platform_multiple_tags(name, platforms, expected_image_tags):
     tags = ["latest", "0.1.0"]
     test_path = f"{TEST_DIR}/test-files/multiplatform"
-    with MultiplatformImageBuilder(cleanup_images=True) as mpib:
+    with MultiplatformImageBuilder() as mpib:
         built_image = mpib.build_multiple_images(
             platforms=platforms,
             path=test_path,
@@ -201,10 +197,6 @@ def test_tag_native_platform_multiple_tags(name, platforms, expected_image_tags)
         )
         assert len(found_image) == 0
 
-    # 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
-
 
 @pytest.mark.parametrize(
     "name, platforms, expected_image_tags",
@@ -214,7 +206,7 @@ def test_tag_native_platform_keep_images(name, platforms, expected_image_tags):
     tag = "latest"
     test_path = f"{TEST_DIR}/test-files/multiplatform"
     try:
-        with MultiplatformImageBuilder(cleanup_images=False) as mpib:
+        with MultiplatformImageBuilder() as mpib:
             built_image = mpib.build_multiple_images(
                 platforms=platforms,
                 path=test_path,
@@ -358,26 +350,25 @@ def test_push_with_dest_names():
     ],
 )
 @patch("buildrunner.docker.multiplatform_image_builder.docker.image.remove")
-@patch("buildrunner.docker.multiplatform_image_builder.docker.image.inspect")
 @patch("buildrunner.docker.multiplatform_image_builder.docker.push")
-@patch("buildrunner.docker.multiplatform_image_builder.docker.image.pull")
+@patch(
+    "buildrunner.docker.multiplatform_image_builder.docker.buildx.imagetools.inspect"
+)
 @patch("buildrunner.docker.multiplatform_image_builder.docker.buildx.build")
 def test_build(
     mock_build,
-    mock_pull,
+    mock_imagetools_inspect,
     mock_push,
-    mock_inspect,
     mock_remove,
     name,
     platforms,
     expected_image_tags,
 ):
     _ = mock_build
-    _ = mock_pull
     _ = mock_push
     _ = mock_remove
-    mock_inspect.return_value = MagicMock()
-    mock_inspect.return_value.id = "myfakeimageid"
+    mock_imagetools_inspect.return_value = MagicMock()
+    mock_imagetools_inspect.return_value.config.digest = "myfakeimageid"
     test_path = f"{TEST_DIR}/test-files/multiplatform"
     with MultiplatformImageBuilder() as mpib:
         built_image = mpib.build_multiple_images(
@@ -397,16 +388,17 @@ def test_build(
 
 
 @patch("buildrunner.docker.multiplatform_image_builder.docker.image.remove")
-@patch("buildrunner.docker.multiplatform_image_builder.docker.image.inspect")
 @patch("buildrunner.docker.multiplatform_image_builder.docker.push")
-@patch("buildrunner.docker.multiplatform_image_builder.docker.image.pull")
+@patch(
+    "buildrunner.docker.multiplatform_image_builder.docker.buildx.imagetools.inspect"
+)
 @patch("buildrunner.docker.multiplatform_image_builder.docker.buildx.build")
 def test_build_multiple_builds(
-    mock_build, mock_pull, mock_push, mock_inspect, mock_remove
+    mock_build, mock_imagetools_inspect, mock_push, mock_remove
 ):
     _ = mock_remove
-    mock_inspect.return_value = MagicMock()
-    mock_inspect.return_value.id = "myfakeimageid"
+    mock_imagetools_inspect.return_value = MagicMock()
+    mock_imagetools_inspect.return_value.config.digest = "myfakeimageid"
     platforms1 = ["linux/amd64", "linux/arm64"]
     expected_image_tags1 = [
         "uuid1-linux-amd64",
@@ -461,11 +453,11 @@ def test_build_multiple_builds(
     prefix = mock_build.call_args.kwargs["tags"][0].split("buildrunner-mp")[0]
     assert mock_build.call_args_list == [
         call(
-            "tests/test-files/multiplatform",
+            test_path,
             tags=[f"{prefix}buildrunner-mp:uuid1-linux-amd64"],
             platforms=["linux/amd64"],
             load=True,
-            file="tests/test-files/multiplatform/Dockerfile",
+            file=f"{test_path}/Dockerfile",
             build_args={"DOCKER_REGISTRY": None},
             builder=None,
             cache=False,
@@ -474,11 +466,11 @@ def test_build_multiple_builds(
             pull=False,
         ),
         call(
-            "tests/test-files/multiplatform",
+            test_path,
             tags=[f"{prefix}buildrunner-mp:uuid1-linux-arm64"],
             platforms=["linux/arm64"],
             load=True,
-            file="tests/test-files/multiplatform/Dockerfile",
+            file=f"{test_path}/Dockerfile",
             build_args={"DOCKER_REGISTRY": None},
             builder=None,
             cache=False,
@@ -487,11 +479,11 @@ def test_build_multiple_builds(
             pull=False,
         ),
         call(
-            "tests/test-files/multiplatform",
+            test_path,
             tags=[f"{prefix}buildrunner-mp:uuid2-linux-amd64"],
             platforms=["linux/amd64"],
             load=True,
-            file="tests/test-files/multiplatform/Dockerfile",
+            file=f"{test_path}/Dockerfile",
             build_args={"DOCKER_REGISTRY": None},
             builder=None,
             cache=False,
@@ -500,11 +492,11 @@ def test_build_multiple_builds(
             pull=False,
         ),
         call(
-            "tests/test-files/multiplatform",
+            test_path,
             tags=[f"{prefix}buildrunner-mp:uuid2-linux-arm64"],
             platforms=["linux/arm64"],
             load=True,
-            file="tests/test-files/multiplatform/Dockerfile",
+            file=f"{test_path}/Dockerfile",
             build_args={"DOCKER_REGISTRY": None},
             builder=None,
             cache=False,
@@ -520,8 +512,8 @@ def test_build_multiple_builds(
         call([f"{prefix}buildrunner-mp:uuid2-linux-amd64"]),
         call([f"{prefix}buildrunner-mp:uuid2-linux-arm64"]),
     ]
-    assert mock_pull.call_count == 4
-    assert mock_pull.call_args_list == [
+    assert mock_imagetools_inspect.call_count == 4
+    assert mock_imagetools_inspect.call_args_list == [
         call(f"{prefix}buildrunner-mp:uuid1-linux-amd64"),
         call(f"{prefix}buildrunner-mp:uuid1-linux-arm64"),
         call(f"{prefix}buildrunner-mp:uuid2-linux-amd64"),

From a4a28027190bf1a5525cce6e8cdd4eed8945e41e Mon Sep 17 00:00:00 2001
From: saville <saville@adobe.com>
Date: Mon, 8 Jan 2024 09:16:58 -0700
Subject: [PATCH 3/3] Code review changes

---
 buildrunner/docker/image_info.py                  | 4 +---
 buildrunner/docker/multiplatform_image_builder.py | 2 +-
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/buildrunner/docker/image_info.py b/buildrunner/docker/image_info.py
index ccdcc88e..d82bc5db 100644
--- a/buildrunner/docker/image_info.py
+++ b/buildrunner/docker/image_info.py
@@ -81,9 +81,7 @@ def built_images(self) -> List[BuiltTaggedImage]:
 
     def image_for_platform(self, platform: str) -> Optional[BuiltTaggedImage]:
         """Retrieves the built image for the given platform."""
-        if platform not in self.images_by_platform:
-            return None
-        return self.images_by_platform[platform]
+        return self.images_by_platform.get(platform)
 
     def add_platform_image(self, platform: str, image_info: BuiltTaggedImage) -> None:
         if platform in self.images_by_platform:
diff --git a/buildrunner/docker/multiplatform_image_builder.py b/buildrunner/docker/multiplatform_image_builder.py
index 44a259a3..966d010e 100644
--- a/buildrunner/docker/multiplatform_image_builder.py
+++ b/buildrunner/docker/multiplatform_image_builder.py
@@ -326,7 +326,7 @@ def _get_single_platform_to_build(self, platforms: List[str]) -> str:
         host_system = python_platform.system()
         host_machine = python_platform.machine()
 
-        if host_system in ("Darwin", "linux"):
+        if host_system.lower() in ("darwin", "linux"):
             native_platform = f"linux/{host_machine}"
         else:
             native_platform = f"{host_system}/{host_machine}"