test: Add vm test controllers.

TEST=./run_tests
BUG=chromium:958030

Change-Id: Ic2469d62015817ea6110a7e7b061d218ec416aa9
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/1592491
Reviewed-by: Evan Hernandez <evanhernandez@chromium.org>
Commit-Queue: Evan Hernandez <evanhernandez@chromium.org>
Tested-by: Evan Hernandez <evanhernandez@chromium.org>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
diff --git a/api/controller/test_unittest.py b/api/controller/test_unittest.py
index 2f23c0a..77faaf2 100644
--- a/api/controller/test_unittest.py
+++ b/api/controller/test_unittest.py
@@ -8,6 +8,7 @@
 from __future__ import print_function
 
 from chromite.api.controller import test as test_controller
+from chromite.api.gen.chromiumos import common_pb2
 from chromite.api.gen.chromite.api import test_pb2
 from chromite.cbuildbot import commands
 from chromite.lib import cros_build_lib
@@ -17,7 +18,7 @@
 from chromite.lib import portage_util
 
 
-class UnitTestTest(cros_test_lib.MockTempDirTestCase):
+class BuildTargetUnitTestTest(cros_test_lib.MockTempDirTestCase):
   """Tests for the UnitTest function."""
 
   def _GetInput(self, board=None, result_path=None, chroot_path=None,
@@ -124,3 +125,75 @@
 
     self.assertNotEqual(0, rc)
     self.assertFalse(output_msg.failed_packages)
+
+
+class VmTestTest(cros_test_lib.MockTestCase):
+  """Test the VmTest endpoint."""
+
+  def _GetInput(self, **kwargs):
+    values = dict(
+        build_target=common_pb2.BuildTarget(name='target'),
+        vm_image=test_pb2.VmTestRequest.VmImage(path='/path/to/image.bin'),
+        test_harness=test_pb2.VmTestRequest.TAST,
+        vm_tests=[test_pb2.VmTestRequest.VmTest(pattern='suite')],
+        ssh_options=test_pb2.VmTestRequest.SshOptions(
+            port=1234, private_key_path='/path/to/id_rsa'),
+    )
+    values.update(kwargs)
+    return test_pb2.VmTestRequest(**values)
+
+  def setUp(self):
+    self.rc_mock = cros_test_lib.RunCommandMock()
+    self.rc_mock.SetDefaultCmdResult()
+    self.StartPatcher(self.rc_mock)
+
+  def testTastAllOptions(self):
+    """Test VmTest for Tast with all options set."""
+    test_controller.VmTest(self._GetInput(), None)
+    self.rc_mock.assertCommandContains([
+        'cros_run_vm_test', '--debug', '--no-display', '--copy-on-write',
+        '--board', 'target',
+        '--image-path', '/path/to/image.bin',
+        '--tast', 'suite',
+        '--ssh-port', '1234',
+        '--private-key', '/path/to/id_rsa',
+    ])
+
+  def testAutotestAllOptions(self):
+    """Test VmTest for Autotest with all options set."""
+    input_proto = self._GetInput(test_harness=test_pb2.VmTestRequest.AUTOTEST)
+    test_controller.VmTest(input_proto, None)
+    self.rc_mock.assertCommandContains([
+        'cros_run_vm_test', '--debug', '--no-display', '--copy-on-write',
+        '--board', 'target',
+        '--image-path', '/path/to/image.bin',
+        '--autotest', 'suite',
+        '--ssh-port', '1234',
+        '--private-key', '/path/to/id_rsa',
+        '--test_that-args=--whitelist-chrome-crashes',
+    ])
+
+  def testMissingBuildTarget(self):
+    """Test VmTest dies when build_target not set."""
+    input_proto = self._GetInput(build_target=None)
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      test_controller.VmTest(input_proto, None)
+
+  def testMissingVmImage(self):
+    """Test VmTest dies when vm_image not set."""
+    input_proto = self._GetInput(vm_image=None)
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      test_controller.VmTest(input_proto, None)
+
+  def testMissingTestHarness(self):
+    """Test VmTest dies when test_harness not specified."""
+    input_proto = self._GetInput(
+        test_harness=test_pb2.VmTestRequest.UNSPECIFIED)
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      test_controller.VmTest(input_proto, None)
+
+  def testMissingVmTests(self):
+    """Test VmTest dies when vm_tests not set."""
+    input_proto = self._GetInput(vm_tests=[])
+    with self.assertRaises(cros_build_lib.DieSystemExit):
+      test_controller.VmTest(input_proto, None)