Take a Swing at Unit Tests

See: aws#6909

samcli/lib/build/app_builder.py

@@ -463,14 +463,14 @@ def _stream_lambda_image_build_logs(self, build_logs: List[Dict[str, str]], func
             raise DockerBuildFailed(msg=f"{function_name} failed to build: {str(ex)}") from ex
 
     def _load_lambda_image(self, image_archive_path: str) -> str:
-        with open(image_archive_path, mode="rb") as image_archive:
-            try:
+        try:
+            with open(image_archive_path, mode="rb") as image_archive:
                 [image, *rest] = self._docker_client.images.load(image_archive)
                 if len(rest) != 0:
                     raise DockerBuildFailed("Archive must represent a single image")
                 return f"{image.id}"
-            except docker.errors.APIError as ex:
-                raise DockerBuildFailed(msg=str(ex)) from ex
+        except (docker.errors.APIError, OSError) as ex:
+            raise DockerBuildFailed(msg=str(ex)) from ex
 
     def _build_layer(
         self,

tests/unit/lib/build_module/test_app_builder.py

@@ -1,12 +1,13 @@
 import os
 import posixpath
 import sys
+from io import BufferedReader
 
 import docker
 import json
 
 from unittest import TestCase
-from unittest.mock import Mock, MagicMock, call, patch, ANY
+from unittest.mock import Mock, MagicMock, call, mock_open, patch, ANY
 from pathlib import Path, WindowsPath
 
 from parameterized import parameterized
@@ -1734,10 +1735,61 @@ def test_can_raise_build_error(self):
             self.builder._build_lambda_image("Name", {}, X86_64)
 
 
+class TestApplicationBuilder_load_lambda_image_function(TestCase):
+    def setUp(self):
+        self.docker_client_mock = Mock()
+        self.builder = ApplicationBuilder(
+            Mock(),
+            "/build/dir",
+            "/base/dir",
+            "/cached/dir",
+            stream_writer=Mock(),
+            docker_client=self.docker_client_mock,
+        )
+
+    @patch("builtins.open", new_callable=mock_open)
+    def test_loads_image_archive(self, mock_open):
+        id = "sha256:1a2b3c4d5e6f"
+
+        self.docker_client_mock.images.load.return_value = [Mock(id=id)]
+
+        image = self.builder._load_lambda_image("./path/to/archive.tar.gz")
+        self.assertEqual(id, image)
+
+    @patch("builtins.open", new_callable=mock_open)
+    def test_must_represent_a_single_image(self, mock_open):
+        self.docker_client_mock.images.load.return_value = [
+            Mock(id="sha256:1a2b3c4d5e6f"),
+            Mock(id="sha256:1f2e3d4c5b6a"),
+        ]
+
+        with self.assertRaises(DockerBuildFailed) as ex:
+            self.builder._load_lambda_image("./path/to/archive.tar.gz")
+        self.assertIn("single", str(ex.exception))
+
+    @patch("builtins.open", side_effect=OSError)
+    def test_image_archive_does_not_exist(self, mock_open):
+        with self.assertRaises(DockerBuildFailed):
+            self.builder._load_lambda_image("./path/to/nowhere.tar.gz")
+
+    @patch("builtins.open", new_callable=mock_open)
+    def test_docker_api_error(self, mock_open):
+        self.docker_client_mock.images.load.side_effect = docker.errors.APIError("failed to dial")
+
+        with self.assertRaises(DockerBuildFailed):
+            self.builder._load_lambda_image("./path/to/archive.tar.gz")
+
+
 class TestApplicationBuilder_build_function(TestCase):
     def setUp(self):
+        self.docker_client_mock = Mock()
         self.builder = ApplicationBuilder(
-            Mock(), "/build/dir", "/base/dir", "cachedir", stream_writer=StreamWriter(sys.stderr)
+            Mock(),
+            "/build/dir",
+            "/base/dir",
+            "cachedir",
+            stream_writer=StreamWriter(sys.stderr),
+            docker_client=self.docker_client_mock,
         )
 
     @patch("samcli.lib.build.app_builder.get_workflow_config")
@@ -2547,6 +2599,19 @@ def test_must_build_in_container_with_custom_default_build_image(self, osutils_m
             specified_workflow=None,
         )
 
+    @parameterized.expand([X86_64, ARM64])
+    @patch.object(Path, "is_file", return_value=True)
+    @patch("builtins.open", new_callable=mock_open)
+    def test_loads_if_path_exists(self, mock_open, mock_is_file, architecture):
+        id = "sha256:1a2b3c4d5e6f"
+        function_name = "function_name"
+        imageuri = str(Path("./path/to/archive.tar.gz"))
+
+        self.docker_client_mock.images.load.return_value = [Mock(id=id)]
+
+        image = self.builder._build_function(function_name, None, imageuri, IMAGE, None, architecture, None, None)
+        self.assertEqual(id, image)
+
 
 class TestApplicationBuilder_build_function_in_process(TestCase):
     def setUp(self):

tests/unit/lib/package/test_ecr_uploader.py

@@ -1,6 +1,7 @@
 from unittest import TestCase
-from unittest.mock import MagicMock, patch, call
+from unittest.mock import MagicMock, Mock, call, mock_open, patch
 
+from pathlib import Path
 from botocore.exceptions import ClientError
 from docker.errors import APIError, BuildError
 from parameterized import parameterized
@@ -194,6 +195,99 @@ def test_upload_failure_while_streaming(self):
         with self.assertRaises(DockerPushFailedError):
             ecr_uploader.upload(image, resource_name="HelloWorldFunction")
 
+    @patch.object(Path, "is_file", return_value=True)
+    @patch("builtins.open", new_callable=mock_open)
+    def test_upload_from_image_archive(self, mock_open, mock_is_file):
+        resource_name = "HelloWorldFunction"
+        digest = "1a2b3c4d5e6f"
+        id = f"sha256:{digest}"
+        image = "./path/to/archive.tar.gz"
+
+        self.docker_client.images.load.return_value = [Mock(id=id)]
+        self.docker_client.api.push.return_value.__iter__.return_value = iter(
+            [
+                {"status": "Pushing to xyz"},
+                {"id": "1", "status": "Preparing", "progress": ""},
+                {"id": "2", "status": "Preparing", "progress": ""},
+                {"id": "3", "status": "Preparing", "progress": ""},
+                {"id": "1", "status": "Pushing", "progress": "[====>   ]"},
+                {"id": "3", "status": "Pushing", "progress": "[====>   ]"},
+                {"id": "2", "status": "Pushing", "progress": "[====>   ]"},
+                {"id": "3", "status": "Pushed", "progress": "[========>]"},
+                {"id": "1", "status": "Pushed", "progress": "[========>]"},
+                {"id": "2", "status": "Pushed", "progress": "[========>]"},
+                {"status": f"image {resource_name} pushed digest: {digest}"},
+                {},
+            ]
+        )
+
+        ecr_uploader = ECRUploader(
+            docker_client=self.docker_client,
+            ecr_client=self.ecr_client,
+            ecr_repo=self.ecr_repo,
+            ecr_repo_multi=self.ecr_repo_multi,
+            tag=self.tag,
+        )
+        ecr_uploader.login = MagicMock()
+        tag = ecr_uploader.upload(image, resource_name=resource_name)
+        self.assertEqual(f"{self.ecr_repo}:{resource_name}-{digest}-{self.tag}", tag)
+
+    @patch.object(Path, "is_file", return_value=False)
+    def test_upload_from_digest(self, mock_is_file):
+        resource_name = "HelloWorldFunction"
+        digest = "1a2b3c4d5e6f"
+        id = f"sha256:{digest}"
+        image = id
+
+        self.docker_client.images.get.return_value = Mock(id=id)
+        self.docker_client.api.push.return_value.__iter__.return_value = iter(
+            [
+                {"status": "Pushing to xyz"},
+                {"id": "1", "status": "Preparing", "progress": ""},
+                {"id": "2", "status": "Preparing", "progress": ""},
+                {"id": "3", "status": "Preparing", "progress": ""},
+                {"id": "1", "status": "Pushing", "progress": "[====>   ]"},
+                {"id": "3", "status": "Pushing", "progress": "[====>   ]"},
+                {"id": "2", "status": "Pushing", "progress": "[====>   ]"},
+                {"id": "3", "status": "Pushed", "progress": "[========>]"},
+                {"id": "1", "status": "Pushed", "progress": "[========>]"},
+                {"id": "2", "status": "Pushed", "progress": "[========>]"},
+                {"status": f"image {resource_name} pushed digest: {digest}"},
+                {},
+            ]
+        )
+
+        ecr_uploader = ECRUploader(
+            docker_client=self.docker_client,
+            ecr_client=self.ecr_client,
+            ecr_repo=self.ecr_repo,
+            ecr_repo_multi=self.ecr_repo_multi,
+            tag=self.tag,
+        )
+        ecr_uploader.login = MagicMock()
+        tag = ecr_uploader.upload(image, resource_name=resource_name)
+        self.assertEqual(f"{self.ecr_repo}:{resource_name}-{digest}-{self.tag}", tag)
+
+    @patch.object(Path, "is_file", return_value=True)
+    @patch("builtins.open", new_callable=mock_open)
+    def test_upload_failure_if_archive_represents_multiple_images(self, mock_open, mock_is_file):
+        resource_name = "HelloWorldFunction"
+        image = "./path/to/archive.tar.gz"
+
+        self.docker_client.images.load.return_value = [Mock(), Mock()]
+
+        ecr_uploader = ECRUploader(
+            docker_client=self.docker_client,
+            ecr_client=self.ecr_client,
+            ecr_repo=self.ecr_repo,
+            ecr_repo_multi=self.ecr_repo_multi,
+            tag=self.tag,
+        )
+        ecr_uploader.login = MagicMock()
+
+        with self.assertRaises(DockerPushFailedError):
+            ecr_uploader.upload(image, resource_name=resource_name)
+
     @patch("samcli.lib.package.ecr_uploader.click.echo")
     def test_delete_artifact_successful(self, patched_click_echo):
         ecr_uploader = ECRUploader(