ImageService/Create: Add factory support.

Add support for building a factory image to ImageService/Create.
Also refactored the code to put more responsibility in the service
layer rather than the controller layer, especially WRT creating
the image paths.
Created a class to simplify the parsed image type result.
Reworked the BuildResult to help track status.

BUG=b:194730793
TEST=./run_tests, cq

Change-Id: Ic541c47660552c86739c64d5a479aa7744a7f164
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/3059809
Tested-by: Alex Klein <saklein@chromium.org>
Commit-Queue: George Engelbrecht <engeg@google.com>
Reviewed-by: George Engelbrecht <engeg@google.com>
Reviewed-by: Sergey Frolov <sfrolov@google.com>
diff --git a/api/controller/image.py b/api/controller/image.py
index 5546a51..5cb6260 100644
--- a/api/controller/image.py
+++ b/api/controller/image.py
@@ -7,9 +7,11 @@
 The image related API endpoints should generally be found here.
 """
 
+import copy
 import functools
 import logging
 import os
+from typing import List, NamedTuple, Set
 
 from chromite.api import controller
 from chromite.api import faux
@@ -79,6 +81,38 @@
     common_pb2.IMAGE_TYPE_GSC_FIRMWARE: constants.IMAGE_TYPE_GSC_FIRMWARE,
 }
 
+# Built image directory symlink names. These names allow specifying a static
+# location for creation to simplify later archival stages. In practice, this
+# sets the symlink argument to build_packages.
+# Core are the build/dev/test images.
+# Use "latest" until we do a better job of passing through image directories,
+# e.g. for artifacts.
+LOCATION_CORE = 'latest'
+# The factory_install image.
+LOCATION_FACTORY = 'factory_shim'
+
+
+class ImageTypes(NamedTuple):
+  """Parsed image types."""
+  images: Set[str]
+  vms: Set[int]
+  mod_images: Set[int]
+
+  @property
+  def core_images(self) -> List[str]:
+    """The core images (base/dev/test) as a list."""
+    return list(self.images - {_IMAGE_MAPPING[_FACTORY_ID]}) or []
+
+  @property
+  def has_factory(self) -> bool:
+    """Whether the factory image is present."""
+    return _IMAGE_MAPPING[_FACTORY_ID] in self.images
+
+  @property
+  def factory(self) -> List[str]:
+    """A list with the factory type if set."""
+    return [_IMAGE_MAPPING[_FACTORY_ID]] if self.has_factory else []
+
 
 def _add_image_to_proto(output_proto, path, image_type, board):
   """Quick helper function to add a new image to the output proto."""
@@ -129,7 +163,8 @@
 
   generated = []
   dlc_func = functools.partial(image.copy_dlc_image, base_path)
-  license_func = functools.partial(image.copy_license_credits, board)
+  license_func = functools.partial(
+      image.copy_license_credits, board, symlink=LOCATION_CORE)
   artifact_types = {
       in_proto.ArtifactType.DLC_IMAGE: dlc_func,
       in_proto.ArtifactType.LICENSE_CREDITS: license_func,
@@ -171,39 +206,67 @@
   # Build the base image if no images provided.
   to_build = input_proto.image_types or [_BASE_ID]
 
-  image_types, vm_types, mod_image_types = _ParseImagesToCreate(to_build)
+  image_types = _ParseImagesToCreate(to_build)
   build_config = _ParseCreateBuildConfig(input_proto)
+  factory_build_config = copy.copy(build_config)
+  build_config.symlink = LOCATION_CORE
+  factory_build_config.symlink = LOCATION_FACTORY
 
+  # Try building the core and factory images.
   # Sorted isn't really necessary here, but it's much easier to test.
-  result = image.Build(
-      board=board, images=sorted(list(image_types)), config=build_config)
+  core_result = image.Build(
+      board, sorted(image_types.core_images), config=build_config)
+  logging.debug('Core Result Images: %s', core_result.images)
 
-  output_proto.success = result.success
+  factory_result = image.Build(
+      board, image_types.factory, config=factory_build_config)
+  logging.debug('Factory Result Images: %s', factory_result.images)
 
-  if result.success:
-    # Success -- we need to list out the images we built in the output.
-    _PopulateBuiltImages(board, image_types, output_proto)
+  # A successful run will have no images missing, will have run at least one
+  # of the two image sets, and neither attempt errored. The no error condition
+  # should be redundant with no missing images, but is cheap insurance.
+  all_built = core_result.all_built and factory_result.all_built
+  one_ran = core_result.build_run or factory_result.build_run
+  no_errors = not core_result.run_error and not factory_result.run_error
+  output_proto.success = success = all_built and one_ran and no_errors
 
-    for vm_type in vm_types:
+  if success:
+    # Success! We need to record the images we built in the output.
+    all_images = {**core_result.images, **factory_result.images}
+    for img_name, img_path in all_images.items():
+      _add_image_to_proto(output_proto, str(img_path), _IMAGE_MAPPING[img_name],
+                          board)
+
+    # Build and record VMs as necessary.
+    for vm_type in image_types.vms:
       is_test = vm_type in [_TEST_VM_ID, _TEST_GUEST_VM_ID]
+      img_type = _IMAGE_MAPPING[_TEST_ID if is_test else _BASE_ID]
+      img_dir = core_result.images[img_type].parent.resolve()
       try:
         if vm_type in [_BASE_GUEST_VM_ID, _TEST_GUEST_VM_ID]:
-          vm_path = image.CreateGuestVm(board, is_test=is_test)
+          vm_path = image.CreateGuestVm(
+              board, is_test=is_test, image_dir=img_dir)
         else:
           vm_path = image.CreateVm(
-              board, disk_layout=build_config.disk_layout, is_test=is_test)
+              board,
+              disk_layout=build_config.disk_layout,
+              is_test=is_test,
+              image_dir=img_dir)
       except image.ImageToVmError as e:
         cros_build_lib.Die(e)
 
       _add_image_to_proto(output_proto, vm_path, vm_type, board)
 
-    for mod_type in mod_image_types:
+    # Build and record any mod images.
+    for mod_type in image_types.mod_images:
       if mod_type == _RECOVERY_ID:
-        base_image_path = _GetBaseImagePath(output_proto)
+        base_image_path = core_result.images[constants.IMAGE_TYPE_BASE]
         result = image.BuildRecoveryImage(
             board=board, image_path=base_image_path)
-        if result.success:
-          _PopulateBuiltImages(board, [_IMAGE_MAPPING[mod_type]], output_proto)
+        if result.all_built:
+          _add_image_to_proto(output_proto,
+                              result.images[_IMAGE_MAPPING[mod_type]], mod_type,
+                              board)
         else:
           cros_build_lib.Die('Failed to create recovery image.')
       else:
@@ -215,36 +278,28 @@
 
   else:
     # Failure, include all of the failed packages in the output when available.
-    if not result.failed_packages:
+    packages = core_result.failed_packages + factory_result.failed_packages
+    if not packages:
       return controller.RETURN_CODE_COMPLETED_UNSUCCESSFULLY
 
-    for package in result.failed_packages:
+    for package in packages:
       current = output_proto.failed_packages.add()
       controller_util.serialize_package_info(package, current)
 
     return controller.RETURN_CODE_UNSUCCESSFUL_RESPONSE_AVAILABLE
 
 
-def _GetBaseImagePath(output_proto):
-  """From image_pb2.CreateImageResult, return base image path or None."""
-  ret = None
-  for i in output_proto.images:
-    if i.type == _BASE_ID:
-      ret = i.path
-  return ret
-
-
-def _ParseImagesToCreate(to_build):
+def _ParseImagesToCreate(to_build: List[int]) -> ImageTypes:
   """Helper function to parse the image types to build.
 
   This function expresses the dependencies of each image type and adds
   the requisite image types if they're not explicitly defined.
 
   Args:
-    to_build (list[int]): The image type list.
+    to_build: The image type list.
 
   Returns:
-    (set, set, set): The image, vm, and mod_image types that need to be built.
+    ImageTypes: The parsed images to build.
   """
   image_types = set()
   vm_types = set()
@@ -271,7 +326,8 @@
   if vm_types.issuperset({_BASE_VM_ID, _TEST_VM_ID}):
     cros_build_lib.Die('Cannot create more than one VM.')
 
-  return image_types, vm_types, mod_image_types
+  return ImageTypes(
+      images=image_types, vms=vm_types, mod_images=mod_image_types)
 
 
 def _ParseCreateBuildConfig(input_proto):
@@ -289,20 +345,6 @@
   )
 
 
-def _PopulateBuiltImages(board, image_types, output_proto):
-  """Helper to list out built images for Create."""
-  # Build out the ImageType->ImagePath mapping in the output.
-  # We're using the default path, so just fetch that, but read the symlink so
-  # the path we're returning is somewhat more permanent.
-  latest_link = image_lib.GetLatestImageLink(board)
-  base_path = os.path.realpath(latest_link)
-
-  for current in image_types:
-    type_id = _IMAGE_MAPPING[current]
-    path = os.path.join(base_path, constants.IMAGE_TYPE_TO_NAME[current])
-    _add_image_to_proto(output_proto, path, type_id, board)
-
-
 def _SignerTestResponse(_input_proto, output_proto, _config):
   """Set output_proto success field on a successful SignerTest response."""
   output_proto.success = True