Add a GCE au_worker

Add a gce_au_worker that runs VMTests on GCE instances.
Extend cros_au_test_harness.py and ctest.py to accept a DUT type of 'gce'.
Add unit tests.

gce_au_worker mimics vm_au_worker, but creates a GCE instance at every call of
UpdateImage. Old images/instances will be deleted at the end of the test. And
like real_au_worker, it calls directly to 'test_that' for the actual autotest
interaction, without any interim external scripts.

This CL uses Lakitu's GCE private project and it should be replaced by one
that's managed by ChromeOS Infrastructure team.

This CL includes an end-to-end test that runs SimpleTestVerify and
SimpleTestUpdateAndVerify in parallel with gce_au_worker. It servces as a demo
but includes some private information that should be stripped before submission.
The tests for update with payloads are coming in a follow up CL.

BUG=brillo:1196
BUG=brillo:1197
TEST=cros_au_test_harness_unittest.py

Change-Id: I5c1104ec28a124c3d85bc7b62a1686a401fb2e95
Reviewed-on: https://chromium-review.googlesource.com/278012
Reviewed-by: Prathmesh Prabhu <pprabhu@chromium.org>
Commit-Queue: Daniel Wang <wonderfly@google.com>
Tested-by: Daniel Wang <wonderfly@google.com>
diff --git a/au_test_harness/au_test.py b/au_test_harness/au_test.py
index 9dd0c03..a04a6e7 100644
--- a/au_test_harness/au_test.py
+++ b/au_test_harness/au_test.py
@@ -13,6 +13,7 @@
 from chromite.lib import cros_logging as logging
 from chromite.lib import dev_server_wrapper
 from crostestutils.au_test_harness import cros_test_proxy
+from crostestutils.au_test_harness import gce_au_worker
 from crostestutils.au_test_harness import real_au_worker
 from crostestutils.au_test_harness import update_exception
 from crostestutils.au_test_harness import vm_au_worker
@@ -39,6 +40,8 @@
     cls.test_results_root = options.test_results_root
     if options.type == 'vm':
       cls.worker_class = vm_au_worker.VMAUWorker
+    elif  options.type == 'gce':
+      cls.worker_class = gce_au_worker.GCEAUWorker
     else:
       cls.worker_class = real_au_worker.RealAUWorker
 
@@ -106,11 +109,11 @@
 
     # Update to
     self.worker.PerformUpdate(self.target_image_path, base_image_path)
-    self.assertTrue(self.worker.VerifyImage())
+    self.assertTrue(self.worker.VerifyImage(self))
 
     # Update from
     self.worker.PerformUpdate(self.target_image_path, self.target_image_path)
-    self.assertTrue(self.worker.VerifyImage())
+    self.assertTrue(self.worker.VerifyImage(self))
 
   def testUpdateWipeStateful(self):
     """Tests if we can update after cleaning the stateful partition.
@@ -126,12 +129,12 @@
     # Update to
     self.worker.PerformUpdate(self.target_image_path, base_image_path,
                               'clean')
-    self.assertTrue(self.worker.VerifyImage())
+    self.assertTrue(self.worker.VerifyImage(self))
 
     # Update from
     self.worker.PerformUpdate(self.target_image_path, self.target_image_path,
                               'clean')
-    self.assertTrue(self.worker.VerifyImage())
+    self.assertTrue(self.worker.VerifyImage(self))
 
   def testInterruptedUpdate(self):
     """Tests what happens if we interrupt payload delivery 3 times."""
@@ -186,7 +189,7 @@
     self.worker.Initialize(9227)
     target_image_path = self.worker.PrepareBase(self.target_image_path)
     self.worker.PerformUpdate(target_image_path, target_image_path)
-    self.assertTrue(self.worker.VerifyImage())
+    self.assertTrue(self.worker.VerifyImage(self))
 
   def SimpleTestVerify(self):
     """Test that only verifies the target image.
@@ -196,7 +199,7 @@
     """
     self.worker.Initialize(9228)
     self.worker.PrepareBase(self.target_image_path)
-    self.assertTrue(self.worker.VerifyImage())
+    self.assertTrue(self.worker.VerifyImage(self))
 
   # --- DISABLED TESTS ---
 
diff --git a/au_test_harness/cros_au_test_harness.py b/au_test_harness/cros_au_test_harness.py
index dd98d7e..c7d189d 100755
--- a/au_test_harness/cros_au_test_harness.py
+++ b/au_test_harness/cros_au_test_harness.py
@@ -28,6 +28,7 @@
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_logging as logging
 from chromite.lib import dev_server_wrapper
+from chromite.lib import gs
 from chromite.lib import parallel
 from chromite.lib import sudo
 from chromite.lib import timeout_util
@@ -75,8 +76,11 @@
         raise parallel.BackgroundFailure(msg)
 
 
-def _ReadUpdateCache(target_image):
+def _ReadUpdateCache(dut_type, target_image):
   """Reads update cache from generate_test_payloads call."""
+  # TODO(wonderfly): Figure out how to use update cache for GCE images.
+  if dut_type == 'gce':
+    return None
   path_to_dump = os.path.dirname(target_image)
   cache_file = os.path.join(path_to_dump, CACHE_FILE)
 
@@ -120,19 +124,34 @@
     options: Parsed options.
     leftover_args: Args left after parsing.
   """
+
+  _IMAGE_PATH_REQUIREMENT = """
+  For vm and real types, the image should be a local file and for gce, the image
+  path has to be a valid Google Cloud Storage path.
+  """
+
   if leftover_args: parser.error('Found unsupported flags ' + leftover_args)
-  if not options.type in ['real', 'vm']:
+  if not options.type in ['real', 'vm', 'gce']:
     parser.error('Failed to specify valid test type.')
 
-  if not options.target_image or not os.path.isfile(options.target_image):
-    parser.error('Testing requires a valid target image.')
+  def _IsValidImage(image):
+    """Asserts that |image_path| is a valid image file for |options.type|."""
+    if not image:
+      return False
+    return (gs.PathIsGs(image) if options.type == 'gce' else
+            os.path.isfile(image))
+
+  if not _IsValidImage(options.target_image):
+    parser.error('Testing requires a valid target image. Given %s. %s' %
+                 (options.target_image, _IMAGE_PATH_REQUIREMENT))
 
   if not options.base_image:
     logging.info('No base image supplied.  Using target as base image.')
     options.base_image = options.target_image
 
-  if not os.path.isfile(options.base_image):
-    parser.error('Testing requires a valid base image.')
+  if not _IsValidImage(options.base_image):
+    parser.error('Testing requires a valid base image. Given: %s. %s' %
+                 (options.base_image, _IMAGE_PATH_REQUIREMENT))
 
   if options.private_key and not os.path.isfile(options.private_key):
     parser.error('Testing requires a valid path to the private key.')
@@ -179,7 +198,7 @@
                     help='Only runs tests with specific prefix i.e. '
                     'testFullUpdateWipeStateful.')
   parser.add_option('-p', '--type', default='vm',
-                    help='type of test to run: [vm, real]. Default: vm.')
+                    help='type of test to run: [vm, real, gce]. Default: vm.')
   parser.add_option('--verbose', default=True, action='store_true',
                     help='Print out rather than capture output as much as '
                     'possible.')
@@ -197,7 +216,7 @@
   CheckOptions(parser, options, leftover_args)
 
   # Generate cache of updates to use during test harness.
-  update_cache = _ReadUpdateCache(options.target_image)
+  update_cache = _ReadUpdateCache(options.type, options.target_image)
   if not update_cache:
     msg = ('No update cache found. Update testing will not work.  Run '
            ' cros_generate_update_payloads if this was not intended.')
@@ -220,7 +239,7 @@
             log_dir=options.test_results_root)
         my_server.Start()
 
-      if options.type == 'vm' and options.parallel:
+      if options.type == 'vm' or options.type == 'gce' and options.parallel:
         _RunTestsInParallel(options)
       else:
         # TODO(sosa) - Take in a machine pool for a real test.
diff --git a/au_test_harness/cros_au_test_harness_unittest.py b/au_test_harness/cros_au_test_harness_unittest.py
new file mode 100755
index 0000000..9b06e4b
--- /dev/null
+++ b/au_test_harness/cros_au_test_harness_unittest.py
@@ -0,0 +1,98 @@
+#!/usr/bin/python2
+#
+# Copyright 2015 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.
+
+"""Tests module containing unittests for cros_au_test_harness.py.
+
+Instead of calling to functions/methods in cros_au_test_harness, tests defined
+here use its binary version, to mimic the behavior of ctest.
+"""
+
+from __future__ import print_function
+
+import os
+import sys
+import unittest
+import uuid
+
+import constants
+sys.path.append(constants.SOURCE_ROOT)
+
+from chromite.lib import cros_build_lib
+
+
+class CrosAuTestHarnessTest(unittest.TestCase):
+  """Testing the GCE related funcionalities in cros_au_test_harness.py"""
+
+  INVALID_TYPE_ERROR = 'Failed to specify valid test type.'
+  INVALID_IMAGE_PATH = 'Testing requires a valid target image.'
+
+  def testCheckOptionsDisallowsUndefinedType(self):
+    """Verifies that CheckOptions complains about invalid type."""
+    cmd = [os.path.join(constants.CROSUTILS_DIR, 'bin', 'cros_au_test_harness'),
+           '--type=aws'
+          ]
+    with self.assertRaises(cros_build_lib.RunCommandError) as cm:
+      cros_build_lib.RunCommand(cmd)
+
+    self.assertIn(self.INVALID_TYPE_ERROR, cm.exception.result.error)
+
+  def testCheckOptionsRecognizesGceType(self):
+    """Verifies that 'gce' is an allowed type."""
+    cmd = [os.path.join(constants.CROSUTILS_DIR, 'bin', 'cros_au_test_harness'),
+           '--type=gce'
+          ]
+    # We don't have all required flags passed in so still expect an exception,
+    # but it won't be complaining about invalid type.
+    with self.assertRaises(cros_build_lib.RunCommandError) as cm:
+      cros_build_lib.RunCommand(cmd)
+
+    self.assertNotIn(self.INVALID_TYPE_ERROR, cm.exception.result.error)
+
+  def testCheckOptionsRequiresGSPathForGCETests(self):
+    """Tests that CheckOptions requires a valid GS path for GCE tests."""
+    local_path = '/tmp/foo/bar'
+    gs_path = 'gs://foo-bucket/bar.tar.gz'
+    cmd = [os.path.join(constants.CROSUTILS_DIR, 'bin', 'cros_au_test_harness'),
+           '--type=gce',
+           '--target_image=%s' % local_path
+          ]
+    with self.assertRaises(cros_build_lib.RunCommandError) as cm:
+      cros_build_lib.RunCommand(cmd)
+    self.assertIn(self.INVALID_IMAGE_PATH, cm.exception.result.error)
+
+    cmd = [os.path.join(constants.CROSUTILS_DIR, 'bin', 'cros_au_test_harness'),
+           '--type=vm',
+           '--target_image=%s' % gs_path
+          ]
+    with self.assertRaises(cros_build_lib.RunCommandError) as cm:
+      cros_build_lib.RunCommand(cmd)
+    self.assertNotIn(self.INVALID_IMAGE_PATH, cm.exception.result.error)
+
+  @unittest.skip('This test runs but only for demo purposes. Do not check it '
+                 'in as is')
+  def testSimpleTestsOnGCE(self):
+    """Tests that cros_au_test_harness is able to run simple tests on GCE.
+
+    Explicitly triggers SimpleTestVerify and SimpleTestUpdateAndVerify via
+    '--test_prefix'.
+    """
+    board = 'lakitu'
+    gs_path = 'gs://test-images/chromiumos_test_image.tar.gz'
+    test_results_dir = 'chroot/tmp/test_results_%s' % str(uuid.uuid4())
+    cmd = [os.path.join(constants.CROSUTILS_DIR, 'bin', 'cros_au_test_harness'),
+           '--type=gce',
+           '--target_image=%s' % gs_path,
+           '--board=%s' % board,
+           '--test_results=%s' % test_results_dir,
+           '--test_prefix=Simple',
+           '--verify_suite_name=smoke',
+           '--parallel'
+          ]
+    cros_build_lib.RunCommand(cmd)
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/au_test_harness/gce_au_worker.py b/au_test_harness/gce_au_worker.py
new file mode 100644
index 0000000..1cc0ccb
--- /dev/null
+++ b/au_test_harness/gce_au_worker.py
@@ -0,0 +1,129 @@
+# Copyright 2015 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.
+
+"""Module containing class that implements an au_worker for GCE instances."""
+
+from __future__ import print_function
+
+import datetime
+import time
+
+from multiprocessing import Process
+
+from chromite.compute import gcloud
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_logging as logging
+from crostestutils.au_test_harness import au_worker
+from crostestutils.au_test_harness import constants
+from crostestutils.au_test_harness import update_exception
+
+
+class GCEAUWorker(au_worker.AUWorker):
+  """Test harness for updating GCE instances."""
+
+  _INSTANCE_PREFIX = 'test-instance-'
+  _IMAGE_PREFIX = 'test-image-'
+
+  def __init__(self, options, test_results_root, project=constants.GCE_PROJECT,
+               zone=constants.GCE_ZONE):
+    """Processes GCE-specific options."""
+    super(GCEAUWorker, self).__init__(options, test_results_root)
+
+    # Google Cloud project and zone, in which to create the test instance.
+    self.gccontext = gcloud.GCContext(project, zone)
+    self.image = ''
+    self.instance = ''
+    self.instance_ip = ''
+
+    # Background processes that delete throw-away instances.
+    self.bg_delete_processes = []
+
+  def CleanUp(self):
+    """Deletes throw-away instances and images"""
+    logging.info('Waiting for all instances and images to be deleted.')
+
+    def _WaitForBackgroundDeleteProcesses():
+      for p in self.bg_delete_processes:
+        p.join()
+
+    _WaitForBackgroundDeleteProcesses()
+    # Delete the instance/image created by the last call to UpdateImage.
+    self._DeleteExistingInstanceInBackground()
+    _WaitForBackgroundDeleteProcesses()
+    logging.info('All instances/images are deleted.')
+
+  def _DeleteExistingInstanceInBackground(self):
+    """Deletes existing instances if any."""
+
+    def _DeleteInstance():
+      bg_delete = Process(target=self.gccontext.DeleteInstance,
+                          args=(self.instance,), kwargs=dict(quiet=True))
+      bg_delete.start()
+      self.bg_delete_processes.append(bg_delete)
+      self.instance = ''
+      self.instance_ip = ''
+
+    def _DeleteImage():
+      bg_delete = Process(target=self.gccontext.DeleteImage,
+                          args=(self.image,), kwargs=dict(quiet=True))
+      bg_delete.start()
+      self.bg_delete_processes.append(bg_delete)
+      self.image = ''
+
+    if self.instance:
+      logging.info('Existing instance %s found. Deleting...', self.instance)
+      _DeleteInstance()
+      if self.image:
+        _DeleteImage()
+
+  def PrepareBase(self, image_path, signed_base=False):
+    """Auto-update to base image to prepare for test."""
+    return self.PrepareRealBase(image_path, signed_base)
+
+  def UpdateImage(self, image_path, src_image_path='', stateful_change='old',
+                  proxy_port=None, private_key_path=None):
+    """Updates the image on a GCE instance.
+
+    Unlike real_au_worker, this method always creates a new instance. Note that
+    |image_path| has to be a valid Google Cloud Storage url, e.g.,
+    gs://foo-bucket/bar.tar.gz.
+    """
+    self._DeleteExistingInstanceInBackground()
+    ts = datetime.datetime.fromtimestamp(time.time()).strftime(
+        '%Y-%m-%d-%H-%M-%S')
+    image = '%s%s' % (self._IMAGE_PREFIX, ts)
+    instance = '%s%s' % (self._INSTANCE_PREFIX, ts)
+
+    # Create an image from |image_path| and an instance from the image.
+    try:
+      self.gccontext.CreateImage(image, image_path)
+      self.gccontext.CreateInstance(instance, image)
+    except gcloud.GCCommandError as e:
+      raise update_exception.UpdateException(
+          1, 'Update failed. Error: %s' % e.message)
+    self.instance_ip = self.gccontext.GetInstanceIP(instance)
+    self.instance = instance
+    self.image = image
+
+  def VerifyImage(self, unittest, percent_required_to_pass=100, test=''):
+    """Verifies an image using test_that with verification suite."""
+    test_directory, _ = self.GetNextResultsPath('autotest_tests')
+    if not test:
+      test = self.verify_suite
+
+    self.TestInfo('Running test %s to verify image.' % test)
+    result = cros_build_lib.RunCommand(
+        ['test_that',
+         '--no-quickmerge',
+         '--results_dir=%s' % test_directory,
+         self.instance_ip,
+         test
+        ], error_code_ok=True, enter_chroot=True, redirect_stdout=True,
+        cwd=constants.CROSUTILS_DIR)
+    ret = self.AssertEnoughTestsPassed(unittest, result.output,
+                                       percent_required_to_pass)
+    if not ret:
+      self._DeleteExistingInstanceInBackground()
+
+    return ret
diff --git a/au_test_harness/vm_au_worker.py b/au_test_harness/vm_au_worker.py
index 6dac628..c093070 100644
--- a/au_test_harness/vm_au_worker.py
+++ b/au_test_harness/vm_au_worker.py
@@ -146,8 +146,12 @@
         cmd, image_path, src_image_path, proxy_port, private_key_path,
         for_vm=True)
 
+  def VerifyImage(self, unittest, percent_required_to_pass=100, test=''):
+    # VMAUWorker disregards |unittest| and |percent_required_to_pass|.
+    return self._VerifyImage(test)
+
   # pylint: disable-msg=W0221
-  def VerifyImage(self, test=''):
+  def _VerifyImage(self, test=''):
     """Runs vm smoke suite or any single test to verify image.
 
     Returns True upon success.  Prints test output and returns False otherwise.
diff --git a/ctest/ctest.py b/ctest/ctest.py
index 8e8fd46..e7454c4 100755
--- a/ctest/ctest.py
+++ b/ctest/ctest.py
@@ -42,7 +42,7 @@
     sign_payloads: Build some payloads with signed keys.
     target: Target image to test.
     test_results_root: Root directory to store au_test_harness results.
-    type: which test harness to run.  Possible values: real, vm.
+    type: which test harness to run.  Possible values: real, vm, gce.
     whitelist_chrome_crashes: Whether to treat Chrome crashes as non-fatal.
   """
 
@@ -231,7 +231,7 @@
                     help='Root directory to store test results.  Should '
                     'be defined relative to chroot root.')
   parser.add_option('--type', default='vm',
-                    help='type of test to run: [vm, real]. Default: vm.')
+                    help='type of test to run: [vm, real, gce]. Default: vm.')
   parser.add_option('--verbose', default=False, action='store_true',
                     help='Print out added debugging information')
   parser.add_option('--whitelist_chrome_crashes', default=False,
@@ -249,11 +249,14 @@
   if args: parser.error('Extra args found %s.' % args)
   if not options.board: parser.error('Need board for image to compare against.')
   if options.only_verify and options.quick_update: parser.error(
-          'Only one of --only_verify or --quick_update should be specified.')
+      'Only one of --only_verify or --quick_update should be specified.')
 
   # force absolute path for these options, since a chdir occurs deeper in the
   # codebase.
   for x in ('nplus1_archive_dir', 'target_image', 'test_results_root'):
+    if x == 'target_image' and options.type == 'gce':
+      # In this case |target_image| is a Google Storage path.
+      continue
     val = getattr(options, x)
     if val is not None:
       setattr(options, x, os.path.abspath(val))
diff --git a/lib/constants.py b/lib/constants.py
index b977d94..9283902 100644
--- a/lib/constants.py
+++ b/lib/constants.py
@@ -18,3 +18,6 @@
 CROSUTILS_LIB_DIR = os.path.join(CROSUTILS_DIR, 'lib')
 
 MAX_TIMEOUT_SECONDS = 4800
+
+GCE_PROJECT = 'chromiumos-gce-testlab'
+GCE_ZONE = 'us-central1-a'