controller: Implement ArtifactsService.

TEST=./run_tests api/controller/artifacts_unittest
BUG=chromium:905039

Change-Id: I3366642cabe29decc63d13feecb0aa9dbce37d1c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/1547916
Tested-by: Evan Hernandez <evanhernandez@chromium.org>
Reviewed-by: Alex Klein <saklein@chromium.org>
Trybot-Ready: Evan Hernandez <evanhernandez@chromium.org>
diff --git a/api/controller/artifacts.py b/api/controller/artifacts.py
new file mode 100644
index 0000000..5423bc7
--- /dev/null
+++ b/api/controller/artifacts.py
@@ -0,0 +1,168 @@
+# -*- coding: utf-8 -*-
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Implements ArtifactService."""
+
+from __future__ import print_function
+
+import os
+
+from chromite.cbuildbot import commands
+from chromite.lib import constants
+from chromite.lib import cros_build_lib
+from chromite.lib import osutils
+
+
+def _GetTargetWorkingDirectory(build_root, target):
+  """Return the working directory for the given build target.
+
+  See commands.py functions for more information on what this means.
+
+  Args:
+    build_root (str): Root CrOS directory being built.
+    target (str): Name of the build target in question.
+
+  Returns:
+    str: Path to the build target's working directory.
+  """
+  return os.path.join(build_root, 'chroot', 'build', target, 'build')
+
+
+def BundleTestUpdatePayloads(input_proto, output_proto):
+  """Generate minimal update payloads for the build target for testing.
+
+  Args:
+    input_proto (BundleRequest): The input proto.
+    output_proto (BundleRequest): The output proto.
+  """
+  target = input_proto.build_target.name
+  output_dir = input_proto.output_dir
+  build_root = constants.SOURCE_ROOT
+
+  # Use the first available image to create the update payload.
+  img_root = os.path.join(build_root, 'src/build/images', 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_root, constants.IMAGE_TYPE_TO_NAME[img_type])
+    if os.path.exists(img_path):
+      img_paths.append(img_path)
+
+  if not img_paths:
+    cros_build_lib.Die(
+        'Expected to find an image of type among %r for target "%s" '
+        'at path %s.', img_types, target, img_root)
+  img = img_paths[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(saklein): Repalce 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)
+
+
+def BundleAutotestFiles(input_proto, output_proto):
+  """Tar the autotest files for a build target.
+
+  Args:
+    input_proto (BundleRequest): The input proto.
+    output_proto (BundleRequest): The output proto.
+  """
+  target = input_proto.build_target.name
+  output_dir = input_proto.output_dir
+  build_root = constants.SOURCE_ROOT
+  cwd = _GetTargetWorkingDirectory(build_root, target)
+
+  # Note that unlike the functions below, this returns the full path
+  # to *multiple* tarballs.
+  # TODO(saklein): Replace with a chromite/service implementation.
+  archives = commands.BuildAutotestTarballsForHWTest(build_root, cwd,
+                                                     output_dir)
+
+  for archive in archives:
+    output_proto.artifacts.add().path = archive
+
+
+def BundleTastFiles(input_proto, output_proto):
+  """Tar the tast files for a build target.
+
+  Args:
+    input_proto (BundleRequest): The input proto.
+    output_proto (BundleRequest): The output proto.
+  """
+  target = input_proto.build_target.name
+  output_dir = input_proto.output_dir
+  build_root = constants.SOURCE_ROOT
+  cwd = _GetTargetWorkingDirectory(build_root, target)
+
+  # Note that unlike the functions below, this returns the full path
+  # to the tarball.
+  # TODO(saklein): Replace with a chromite/service implementation.
+  archive = commands.BuildTastBundleTarball(build_root, cwd, output_dir)
+
+  output_proto.artifacts.add().path = archive
+
+
+def BundlePinnedGuestImages(input_proto, output_proto):
+  """Tar the pinned guest images for a build target.
+
+  Args:
+    input_proto (BundleRequest): The input proto.
+    output_proto (BundleRequest): The output proto.
+  """
+  target = input_proto.build_target.name
+  output_dir = input_proto.output_dir
+  build_root = constants.SOURCE_ROOT
+
+  # TODO(saklein): Replace with a chromite/service implementation.
+  archive = commands.BuildPinnedGuestImagesTarball(build_root, target,
+                                                   output_dir)
+
+  output_proto.artifacts.add().path = os.path.join(output_dir, archive)
+
+
+def BundleFirmware(input_proto, output_proto):
+  """Tar the firmware images for a build target.
+
+  Args:
+    input_proto (BundleRequest): The input proto.
+    output_proto (BundleRequest): The output proto.
+  """
+  target = input_proto.build_target.name
+  output_dir = input_proto.output_dir
+  build_root = constants.SOURCE_ROOT
+
+  # TODO(saklein): Replace with a chromite/service implementation.
+  archive = commands.BuildFirmwareArchive(build_root, target, output_dir)
+
+  output_proto.artifacts.add().path = os.path.join(output_dir, archive)
+
+
+def BundleEbuildLogs(input_proto, output_proto):
+  """Tar the ebuild logs for a build target.
+
+  Args:
+    input_proto (BundleRequest): The input proto.
+    output_proto (BundleRequest): The output proto.
+  """
+  target = input_proto.build_target.name
+  output_dir = input_proto.output_dir
+  build_root = constants.SOURCE_ROOT
+
+  # TODO(saklein): Replace with a chromite/service implementation.
+  archive = commands.BuildEbuildLogsTarball(build_root, target, output_dir)
+
+  output_proto.artifacts.add().path = os.path.join(output_dir, archive)
diff --git a/api/controller/artifacts_unittest b/api/controller/artifacts_unittest
new file mode 120000
index 0000000..ef3e37b
--- /dev/null
+++ b/api/controller/artifacts_unittest
@@ -0,0 +1 @@
+../../scripts/wrapper.py
\ No newline at end of file
diff --git a/api/controller/artifacts_unittest.py b/api/controller/artifacts_unittest.py
new file mode 100644
index 0000000..b253220
--- /dev/null
+++ b/api/controller/artifacts_unittest.py
@@ -0,0 +1,186 @@
+# -*- coding: utf-8 -*-
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Unittests for Artifacts operations."""
+
+from __future__ import print_function
+
+import mock
+import os
+
+from chromite.api.controller import artifacts
+from chromite.api.gen.chromite.api import artifacts_pb2
+from chromite.cbuildbot import commands
+from chromite.lib import constants
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_test_lib
+from chromite.lib import osutils
+
+
+class BundleTestCase(cros_test_lib.MockTestCase):
+  """Basic setup for all artifacts unittests."""
+
+  def setUp(self):
+    self.input_proto = artifacts_pb2.BundleRequest()
+    self.input_proto.build_target.name = 'target'
+    self.input_proto.output_dir = '/tmp/artifacts'
+    self.output_proto = artifacts_pb2.BundleResponse()
+
+    self.PatchObject(constants, 'SOURCE_ROOT', new='/cros')
+
+
+class BundleAutotestFilesTest(BundleTestCase):
+  """Unittests for BundleAutotestFiles."""
+
+  def testBundleAutotestFiles(self):
+    """BundleAutotestFiles calls cbuildbot/commands with correct args."""
+    build_autotest_tarballs = self.PatchObject(
+        commands,
+        'BuildAutotestTarballsForHWTest',
+        return_value=[
+            '/tmp/artifacts/autotest-a.tar.gz',
+            '/tmp/artifacts/autotest-b.tar.gz',
+        ])
+    artifacts.BundleAutotestFiles(self.input_proto, self.output_proto)
+    self.assertItemsEqual([
+        artifact.path for artifact in self.output_proto.artifacts
+    ], ['/tmp/artifacts/autotest-a.tar.gz', '/tmp/artifacts/autotest-b.tar.gz'])
+    self.assertEqual(build_autotest_tarballs.call_args_list, [
+        mock.call('/cros', '/cros/chroot/build/target/build', '/tmp/artifacts')
+    ])
+
+
+class BundleTastFilesTest(BundleTestCase):
+  """Unittests for BundleTastFiles."""
+
+  def testBundleTastFiles(self):
+    """BundleTastFiles calls cbuildbot/commands with correct args."""
+    build_tast_bundle_tarball = self.PatchObject(
+        commands,
+        'BuildTastBundleTarball',
+        return_value='/tmp/artifacts/tast.tar.gz')
+    artifacts.BundleTastFiles(self.input_proto, self.output_proto)
+    self.assertEqual(
+        [artifact.path for artifact in self.output_proto.artifacts],
+        ['/tmp/artifacts/tast.tar.gz'])
+    self.assertEqual(build_tast_bundle_tarball.call_args_list, [
+        mock.call('/cros', '/cros/chroot/build/target/build', '/tmp/artifacts')
+    ])
+
+
+class BundlePinnedGuestImagesTest(BundleTestCase):
+  """Unittests for BundlePinnedGuestImages."""
+
+  def testBundlePinnedGuestImages(self):
+    """BundlePinnedGuestImages calls cbuildbot/commands with correct args."""
+    build_pinned_guest_images_tarball = self.PatchObject(
+        commands,
+        'BuildPinnedGuestImagesTarball',
+        return_value='pinned-guest-images.tar.gz')
+    artifacts.BundlePinnedGuestImages(self.input_proto, self.output_proto)
+    self.assertEqual(
+        [artifact.path for artifact in self.output_proto.artifacts],
+        ['/tmp/artifacts/pinned-guest-images.tar.gz'])
+    self.assertEqual(build_pinned_guest_images_tarball.call_args_list,
+                     [mock.call('/cros', 'target', '/tmp/artifacts')])
+
+
+class BundleFirmwareTest(BundleTestCase):
+  """Unittests for BundleFirmware."""
+
+  def testBundleFirmware(self):
+    """BundleFirmware calls cbuildbot/commands with correct args."""
+    build_firmware_archive = self.PatchObject(
+        commands, 'BuildFirmwareArchive', return_value='firmware.tar.gz')
+    artifacts.BundleFirmware(self.input_proto, self.output_proto)
+    self.assertEqual(
+        [artifact.path for artifact in self.output_proto.artifacts],
+        ['/tmp/artifacts/firmware.tar.gz'])
+    self.assertEqual(build_firmware_archive.call_args_list,
+                     [mock.call('/cros', 'target', '/tmp/artifacts')])
+
+
+class BundleEbuildLogsTest(BundleTestCase):
+  """Unittests for BundleEbuildLogs."""
+
+  def testBundleEbuildLogs(self):
+    """BundleEbuildLogs calls cbuildbot/commands with correct args."""
+    build_ebuild_logs_tarball = self.PatchObject(
+        commands, 'BuildEbuildLogsTarball', return_value='ebuild-logs.tar.gz')
+    artifacts.BundleEbuildLogs(self.input_proto, self.output_proto)
+    self.assertEqual(
+        [artifact.path for artifact in self.output_proto.artifacts],
+        ['/tmp/artifacts/ebuild-logs.tar.gz'])
+    self.assertEqual(build_ebuild_logs_tarball.call_args_list,
+                     [mock.call('/cros', 'target', '/tmp/artifacts')])
+
+
+class BundleTestUpdatePayloadsTest(cros_test_lib.MockTempDirTestCase):
+  """Unittests for BundleTestUpdatePayloads."""
+
+  def setUp(self):
+    self.source_root = os.path.join(self.tempdir, 'cros')
+    osutils.SafeMakedirs(self.source_root)
+
+    self.archive_root = os.path.join(self.tempdir, 'output')
+    osutils.SafeMakedirs(self.archive_root)
+
+    self.target = 'target'
+    self.image_root = os.path.join(self.source_root, 'src/build/images/target')
+
+    self.input_proto = artifacts_pb2.BundleRequest()
+    self.input_proto.build_target.name = self.target
+    self.input_proto.output_dir = self.archive_root
+    self.output_proto = artifacts_pb2.BundleResponse()
+
+    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)
+
+    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)
+
+  def testBundleTestUpdatePayloads(self):
+    """BundleTestUpdatePayloads calls cbuildbot/commands with correct args."""
+    image_path = os.path.join(self.image_root, constants.BASE_IMAGE_BIN)
+    osutils.WriteFile(image_path, 'image!', makedirs=True)
+
+    artifacts.BundleTestUpdatePayloads(self.input_proto, self.output_proto)
+
+    actual = [
+        os.path.relpath(artifact.path, self.archive_root)
+        for artifact in self.output_proto.artifacts
+    ]
+    expected = ['payload.bin', 'payload-qp.bin']
+    self.assertItemsEqual(actual, expected)
+
+    actual = [
+        os.path.relpath(path, self.archive_root)
+        for path in osutils.DirectoryIterator(self.archive_root)
+    ]
+    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 testBundleTestUpdatePayloadsNoImage(self):
+    """BundleTestUpdatePayloads dies if no usable image is found for target."""
+    # Intentionally do not write image.
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      artifacts.BundleTestUpdatePayloads(self.input_proto, self.output_proto)