Build API: Allow building VM images from the ImageService.Create endpoint.

BUG=chromium:961431
TEST=manual, run_tests

Change-Id: I03b9fb9c00728afc6d29b5dedc5580fd865858fe
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/1603406
Reviewed-by: Alex Klein <saklein@chromium.org>
Tested-by: Alex Klein <saklein@chromium.org>
diff --git a/api/controller/image_unittest.py b/api/controller/image_unittest.py
index f26abb1..d160287 100644
--- a/api/controller/image_unittest.py
+++ b/api/controller/image_unittest.py
@@ -23,6 +23,21 @@
 class CreateTest(cros_test_lib.MockTempDirTestCase):
   """Create image tests."""
 
+  def _GetRequest(self, board=None, types=None, version=None, builder_path=None,
+                  disable_rootfs_verification=False):
+    """Helper to build a request instance."""
+    return image_pb2.CreateImageRequest(
+        build_target={'name': board},
+        image_types=types,
+        disable_rootfs_verification=disable_rootfs_verification,
+        version=version,
+        builder_path=builder_path,
+    )
+
+  def _GetResponse(self):
+    """Helper to build an empty response instance."""
+    return image_pb2.CreateImageResult()
+
   def testArgumentValidation(self):
     """Test the argument validation."""
     input_proto = image_pb2.CreateImageRequest()
@@ -32,32 +47,47 @@
     with self.assertRaises(cros_build_lib.DieSystemExit):
       image_controller.Create(input_proto, output_proto)
 
-  def testImageTypeHandling(self):
-    """Test the image type handling."""
+  def testNoTypeSpecified(self):
+    """Test the image type default."""
+    request = self._GetRequest(board='board')
+    response = self._GetResponse()
+
     # Failed result to avoid the success handling logic.
     result = image_service.BuildResult(1, [])
     build_patch = self.PatchObject(image_service, 'Build', return_value=result)
-    input_proto = image_pb2.CreateImageRequest()
-    input_proto.build_target.name = 'board'
-    output_proto = image_pb2.CreateImageResult()
 
-    # Should default to the base image.
-    image_controller.Create(input_proto, output_proto)
+    image_controller.Create(request, response)
     build_patch.assert_called_with(images=[constants.IMAGE_TYPE_BASE],
-                                   board=u'board', config=mock.ANY)
+                                   board='board', config=mock.ANY)
 
-    # Should be using a value that's provided.
-    input_proto.image_types.append(common_pb2.DEV)
-    image_controller.Create(input_proto, output_proto)
+  def testSingleTypeSpecified(self):
+    """Test it's properly using a specified type."""
+    request = self._GetRequest(board='board', types=[common_pb2.DEV])
+    response = self._GetResponse()
+
+    # Failed result to avoid the success handling logic.
+    result = image_service.BuildResult(1, [])
+    build_patch = self.PatchObject(image_service, 'Build', return_value=result)
+
+    image_controller.Create(request, response)
     build_patch.assert_called_with(images=[constants.IMAGE_TYPE_DEV],
-                                   board=u'board', config=mock.ANY)
+                                   board='board', config=mock.ANY)
 
-    input_proto.image_types.append(common_pb2.BASE)
-    input_proto.image_types.append(common_pb2.TEST)
-    expected_images = [constants.IMAGE_TYPE_BASE, constants.IMAGE_TYPE_DEV,
-                       constants.IMAGE_TYPE_TEST]
-    image_controller.Create(input_proto, output_proto)
-    build_patch.assert_called_with(images=expected_images, board=u'board',
+  def testMultipleAndImpliedTypes(self):
+    """Test multiple types and implied type handling."""
+    # The TEST_VM type should force it to build the test image.
+    types = [common_pb2.BASE, common_pb2.TEST_VM]
+    expected_images = [constants.IMAGE_TYPE_BASE, constants.IMAGE_TYPE_TEST]
+
+    request = self._GetRequest(board='board', types=types)
+    response = self._GetResponse()
+
+    # Failed result to avoid the success handling logic.
+    result = image_service.BuildResult(1, [])
+    build_patch = self.PatchObject(image_service, 'Build', return_value=result)
+
+    image_controller.Create(request, response)
+    build_patch.assert_called_with(images=expected_images, board='board',
                                    config=mock.ANY)
 
   def testFailedPackageHandling(self):