artifacts: Add BundleImageZip endpoint.

TEST=./run_tests
BUG=chromium:905039

Change-Id: I747befa69be76f2bede74456f1ae47144e038a76
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/1559330
Reviewed-by: Sean Abraham <seanabraham@chromium.org>
Reviewed-by: Jason Clinton <jclinton@chromium.org>
Tested-by: Evan Hernandez <evanhernandez@chromium.org>
diff --git a/api/controller/artifacts_unittest.py b/api/controller/artifacts_unittest.py
index 7a3f23e..c2d20cd 100644
--- a/api/controller/artifacts_unittest.py
+++ b/api/controller/artifacts_unittest.py
@@ -31,6 +31,29 @@
     self.PatchObject(constants, 'SOURCE_ROOT', new='/cros')
 
 
+class BundleImageZipTest(BundleTestCase):
+  """Unittests for BundleImageZip."""
+
+  def testBundleImageZip(self):
+    """BundleImageZip calls cbuildbot/commands with correct args."""
+    build_image_zip = self.PatchObject(
+        commands, 'BuildImageZip', return_value='image.zip')
+    self.PatchObject(os.path, 'exists', return_value=True)
+    artifacts.BundleImageZip(self.input_proto, self.output_proto)
+    self.assertEqual(
+        [artifact.path for artifact in self.output_proto.artifacts],
+        ['/tmp/artifacts/image.zip'])
+    self.assertEqual(
+        build_image_zip.call_args_list,
+        [mock.call('/tmp/artifacts', '/cros/src/build/images/target/latest')])
+
+  def testBundleImageZipNoImageDir(self):
+    """BundleImageZip dies when image dir does not exist."""
+    self.PatchObject(os.path, 'exists', return_value=False)
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      artifacts.BundleImageZip(self.input_proto, self.output_proto)
+
+
 class BundleAutotestFilesTest(BundleTestCase):
   """Unittests for BundleAutotestFiles."""
 
@@ -182,8 +205,15 @@
     self.assertEqual(self.generate_quick_provision_payloads.call_args_list,
                      [mock.call(image_path, mock.ANY)])
 
+  def testBundleTestUpdatePayloadsNoImageDir(self):
+    """BundleTestUpdatePayloads dies if no image dir is found."""
+    # Intentionally do not write image directory.
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      artifacts.BundleTestUpdatePayloads(self.input_proto, self.output_proto)
+
   def testBundleTestUpdatePayloadsNoImage(self):
     """BundleTestUpdatePayloads dies if no usable image is found for target."""
-    # Intentionally do not write image.
+    # Intentionally do not write image, but create the directory.
+    osutils.SafeMakedirs(self.image_root)
     with self.assertRaises(cros_build_lib.DieSystemExit):
       artifacts.BundleTestUpdatePayloads(self.input_proto, self.output_proto)