Use custom key for VM tests

BUG=brillo:1200
TEST=Passed trybot run with rambi-pre-cq

Change-Id: Ic822f4639b0f137e512d94fac0f7e60e505dd2b5
Reviewed-on: https://chromium-review.googlesource.com/288817
Commit-Ready: Daniel Wang <wonderfly@google.com>
Tested-by: Daniel Wang <wonderfly@google.com>
Reviewed-by: Prathmesh Prabhu <pprabhu@chromium.org>
diff --git a/au_test_harness/au_worker.py b/au_test_harness/au_worker.py
index 26b7278..7e80cb5 100644
--- a/au_test_harness/au_worker.py
+++ b/au_test_harness/au_worker.py
@@ -15,6 +15,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 path_util
 from crostestutils.au_test_harness import update_exception
 
 
@@ -39,8 +40,11 @@
     else:
       self.verify_suite = 'suite:%s' % (options.verify_suite_name or 'smoke')
 
-    # An optional ssh private key for testing.
+    # An optional ssh private key for testing. This path is going to be used by
+    # test_that, so it has to a valid chroot path.
     self.ssh_private_key = options.ssh_private_key
+    if self.ssh_private_key:
+      self.ssh_private_key = path_util.ToChrootPath(self.ssh_private_key)
 
   def CleanUp(self):
     """Called at the end of every test."""
diff --git a/au_test_harness/vm_au_worker.py b/au_test_harness/vm_au_worker.py
index 7d2427f..4bfa724 100644
--- a/au_test_harness/vm_au_worker.py
+++ b/au_test_harness/vm_au_worker.py
@@ -176,8 +176,8 @@
     if self.whitelist_chrome_crashes:
       command.append('--whitelist_chrome_crashes')
     if self.ssh_private_key is not None:
-      logging.warning('Flag "--ssh_private_key" set but not yet supported for '
-                      '"VMAUWorker". Default test key will be used.')
+      command.append('--ssh_private_key=%s' % self.ssh_private_key)
+
     self.TestInfo('Running smoke suite to verify image.')
     result = cros_build_lib.RunCommand(
         command, print_cmd=False, combine_stdout_stderr=True,
diff --git a/cros_run_vm_test b/cros_run_vm_test
index 3eec0a7..90f19f2 100755
--- a/cros_run_vm_test
+++ b/cros_run_vm_test
@@ -25,6 +25,8 @@
 DEFINE_integer verbose 1 "{0,1,2} Max verbosity shows autoserv debug output." v
 DEFINE_boolean whitelist_chrome_crashes ${FLAGS_FALSE} \
     "Treat Chrome crashes as non-fatal."
+DEFINE_string ssh_private_key "" \
+    "Path to the private key to use to ssh into test image as the root user."
 
 set -e
 
@@ -80,6 +82,13 @@
   test_args+=( "${FLAGS_args}" )
 fi
 
+if [ -n "${FLAGS_ssh_private_key}" ]; then
+  test_args+=(
+    "--ssh_private_key"
+    "${FLAGS_ssh_private_key}"
+  )
+fi
+
 trap stop_kvm EXIT
 start_kvm "${FLAGS_image_path}" "${FLAGS_board}"
 retry_until_ssh