Build API: Add ImageService/CreateVm method.

BUG=chromium:957758, b:131710961
TEST=run_tests
Cq-Depend: chromium:1592723, chromium:1593932

Change-Id: Id2a481d3e22243ba2846279f61d55713b40983d3
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/1593420
Commit-Queue: Alex Klein <saklein@chromium.org>
Tested-by: Alex Klein <saklein@chromium.org>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
Reviewed-by: Evan Hernandez <evanhernandez@chromium.org>
diff --git a/api/controller/image_unittest.py b/api/controller/image_unittest.py
index 36d3766..f35e87a 100644
--- a/api/controller/image_unittest.py
+++ b/api/controller/image_unittest.py
@@ -13,6 +13,7 @@
 from chromite.api.gen.chromite.api import image_pb2
 from chromite.api.controller import image as image_controller
 from chromite.lib import constants
+from chromite.lib import cros_build_lib
 from chromite.lib import cros_test_lib
 from chromite.lib import osutils
 from chromite.service import image as image_service
@@ -27,7 +28,7 @@
     output_proto = image_pb2.CreateImageResult()
 
     # No board should cause it to fail.
-    with self.assertRaises(image_controller.InvalidArgumentError):
+    with self.assertRaises(cros_build_lib.DieSystemExit):
       image_controller.Create(input_proto, output_proto)
 
   def testImageTypeHandling(self):
@@ -73,6 +74,68 @@
       self.assertIn((package.category, package.package_name), expected_packages)
 
 
+class CreateVmTest(cros_test_lib.MockTestCase):
+  """CreateVm tests."""
+
+  def _GetInput(self, board=None, image_type=None):
+    """Helper to create an input proto instance."""
+    # pylint: disable=protected-access
+
+    return image_pb2.CreateVmRequest(
+        image={'build_target': {'name': board}, 'type': image_type})
+
+  def _GetOutput(self):
+    """Helper to create an empty output proto instance."""
+    return image_pb2.CreateVmResponse()
+
+  def testNoArgsFails(self):
+    """Make sure it fails with no arguments."""
+    request = self._GetInput()
+    response = self._GetOutput()
+
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      image_controller.CreateVm(request, response)
+
+  def testNoBuildTargetFails(self):
+    """Make sure it fails with no build target."""
+    request = self._GetInput(image_type=image_pb2.Image.TEST)
+    response = self._GetOutput()
+
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      image_controller.CreateVm(request, response)
+
+  def testNoTypeFails(self):
+    """Make sure it fails with no build target."""
+    request = self._GetInput(board='board')
+    response = self._GetOutput()
+
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      image_controller.CreateVm(request, response)
+
+  def testTestImage(self):
+    """Make sure the test image identification works properly."""
+    request = self._GetInput(board='board', image_type=image_pb2.Image.TEST)
+    response = self._GetOutput()
+    create_patch = self.PatchObject(image_service, 'CreateVm',
+                                    return_value='/vm/path')
+
+    image_controller.CreateVm(request, response)
+
+    create_patch.assert_called_once_with('board', chroot=mock.ANY, is_test=True)
+
+  def testNonTestImage(self):
+    """Make sure the test image identification works properly."""
+    request = self._GetInput(board='board', image_type=image_pb2.Image.BASE)
+    response = self._GetOutput()
+    create_patch = self.PatchObject(image_service, 'CreateVm',
+                                    return_value='/vm/path')
+
+    image_controller.CreateVm(request, response)
+
+    create_patch.assert_called_once_with('board', chroot=mock.ANY,
+                                         is_test=False)
+
+
 class ImageTest(cros_test_lib.MockTempDirTestCase):
   """Image service tests."""
 
@@ -91,22 +154,22 @@
     output_proto = image_pb2.TestImageResult()
 
     # Nothing provided.
-    with self.assertRaises(image_controller.InvalidArgumentError):
+    with self.assertRaises(cros_build_lib.DieSystemExit):
       image_controller.Test(input_proto, output_proto)
 
     # Just one argument.
     input_proto.build_target.name = self.board
-    with self.assertRaises(image_controller.InvalidArgumentError):
+    with self.assertRaises(cros_build_lib.DieSystemExit):
       image_controller.Test(input_proto, output_proto)
 
     # Two arguments provided.
     input_proto.result.directory = self.result_directory
-    with self.assertRaises(image_controller.InvalidArgumentError):
+    with self.assertRaises(cros_build_lib.DieSystemExit):
       image_controller.Test(input_proto, output_proto)
 
     # Invalid image path.
     input_proto.image.path = '/invalid/image/path'
-    with self.assertRaises(image_controller.InvalidArgumentError):
+    with self.assertRaises(cros_build_lib.DieSystemExit):
       image_controller.Test(input_proto, output_proto)
 
     # All valid arguments.