BundleTestUpdatePayloads: Refactor to service.

BUG=chromium:954283
TEST=run_tests, manually ran endpoint.

Change-Id: I78da1164acb3b8d048fd202e4c9e50c1dbd73b43
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/1679378
Tested-by: Alex Klein <saklein@chromium.org>
Auto-Submit: Alex Klein <saklein@chromium.org>
Reviewed-by: Evan Hernandez <evanhernandez@chromium.org>
Commit-Queue: Alex Klein <saklein@chromium.org>
diff --git a/api/controller/artifacts.py b/api/controller/artifacts.py
index 68a9aa8..d7c2b64 100644
--- a/api/controller/artifacts.py
+++ b/api/controller/artifacts.py
@@ -7,6 +7,7 @@
 
 from __future__ import print_function
 
+import functools
 import os
 
 from chromite.api import validate
@@ -18,7 +19,6 @@
 from chromite.lib import constants
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_logging as logging
-from chromite.lib import osutils
 from chromite.lib import sysroot_lib
 from chromite.service import artifacts
 
@@ -72,35 +72,21 @@
 
   # Use the first available image to create the update payload.
   img_dir = _GetImageDir(build_root, target)
-  img_types = [
-      constants.IMAGE_TYPE_TEST, constants.IMAGE_TYPE_DEV,
-      constants.IMAGE_TYPE_BASE
-  ]
-  img_paths = []
-  for img_type in img_types:
-    img_path = os.path.join(img_dir, constants.IMAGE_TYPE_TO_NAME[img_type])
-    if os.path.exists(img_path):
-      img_paths.append(img_path)
+  img_types = [constants.IMAGE_TYPE_TEST, constants.IMAGE_TYPE_DEV,
+               constants.IMAGE_TYPE_BASE]
+  img_names = [constants.IMAGE_TYPE_TO_NAME[t] for t in img_types]
+  img_paths = map(functools.partial(os.path.join, img_dir), img_names)
+  valid_images = filter(os.path.exists, img_paths)
 
-  if not img_paths:
+  if not valid_images:
     cros_build_lib.Die(
         'Expected to find an image of type among %r for target "%s" '
         'at path %s.', img_types, target, img_dir)
-  img = img_paths[0]
+  image = valid_images[0]
 
-  # Unfortunately, the relevant commands.py functions do not return
-  # a list of generated files. As a workaround, we have commands.py
-  # put the files in a separate temporary directory so we can catalog them,
-  # then move them to the output dir.
-  # TODO(crbug.com/954283): Replace with a chromite/service implementation.
-  with osutils.TempDir() as temp:
-    commands.GeneratePayloads(img, temp, full=True, stateful=True, delta=True)
-    commands.GenerateQuickProvisionPayloads(img, temp)
-    for path in osutils.DirectoryIterator(temp):
-      if os.path.isfile(path):
-        rel_path = os.path.relpath(path, temp)
-        output_proto.artifacts.add().path = os.path.join(output_dir, rel_path)
-    osutils.CopyDirContents(temp, output_dir, allow_nonempty=True)
+  payloads = artifacts.BundleTestUpdatePayloads(image, output_dir)
+  for payload in payloads:
+    output_proto.artifacts.add().path = payload
 
 
 def BundleAutotestFiles(input_proto, output_proto):
diff --git a/api/controller/artifacts_unittest.py b/api/controller/artifacts_unittest.py
index f0740a1..d4e5631 100644
--- a/api/controller/artifacts_unittest.py
+++ b/api/controller/artifacts_unittest.py
@@ -305,20 +305,14 @@
 
     self.PatchObject(constants, 'SOURCE_ROOT', new=self.source_root)
 
-    def MockGeneratePayloads(image_path, archive_dir, **kwargs):
-      assert kwargs
-      osutils.WriteFile(os.path.join(archive_dir, 'payload.bin'), image_path)
+    def MockPayloads(image_path, archive_dir):
+      osutils.WriteFile(os.path.join(archive_dir, 'payload1.bin'), image_path)
+      osutils.WriteFile(os.path.join(archive_dir, 'payload2.bin'), image_path)
+      return [os.path.join(archive_dir, 'payload1.bin'),
+              os.path.join(archive_dir, 'payload2.bin')]
 
-    self.generate_payloads = self.PatchObject(
-        commands, 'GeneratePayloads', side_effect=MockGeneratePayloads)
-
-    def MockGenerateQuickProvisionPayloads(image_path, archive_dir):
-      osutils.WriteFile(os.path.join(archive_dir, 'payload-qp.bin'), image_path)
-
-    self.generate_quick_provision_payloads = self.PatchObject(
-        commands,
-        'GenerateQuickProvisionPayloads',
-        side_effect=MockGenerateQuickProvisionPayloads)
+    self.bundle_patch = self.PatchObject(
+        artifacts_svc, 'BundleTestUpdatePayloads', side_effect=MockPayloads)
 
   def testBundleTestUpdatePayloads(self):
     """BundleTestUpdatePayloads calls cbuildbot/commands with correct args."""
@@ -331,7 +325,7 @@
         os.path.relpath(artifact.path, self.archive_root)
         for artifact in self.output_proto.artifacts
     ]
-    expected = ['payload.bin', 'payload-qp.bin']
+    expected = ['payload1.bin', 'payload2.bin']
     self.assertItemsEqual(actual, expected)
 
     actual = [
@@ -340,13 +334,6 @@
     ]
     self.assertItemsEqual(actual, expected)
 
-    self.assertEqual(self.generate_payloads.call_args_list, [
-        mock.call(image_path, mock.ANY, full=True, stateful=True, delta=True),
-    ])
-
-    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.