au_worker: Defer chroot path replacement until call to test_that

When running KVM tests, cros_run_vm_test starts a KVM instances, and tries to
SSH into it with the provided private key, before its call to test_that (the
ssh_ping phase). ssh_ping needes the key path to be outside of chroot. So we
defer the chroot path replacement until the actual call to test_that, which only
runs inside chroot.

Also, fix an out-dated parameter in gce_au_worker.py. This is a left-over of
CL:286934.

BUG=brillo:1200
TEST=Trybot run with rambi-pre-cq, and a trybot run with a build that has the
USE flag 'generated_key' enabled.
CQ-DEPEND=CL:296723

Change-Id: Iaf911d859333d748add452bda1a2d012f91cf2f8
Reviewed-on: https://chromium-review.googlesource.com/296832
Commit-Ready: Daniel Wang <wonderfly@google.com>
Tested-by: Daniel Wang <wonderfly@google.com>
Reviewed-by: Don Garrett <dgarrett@chromium.org>
diff --git a/au_test_harness/au_worker.py b/au_test_harness/au_worker.py
index 7e80cb5..b2062cb 100644
--- a/au_test_harness/au_worker.py
+++ b/au_test_harness/au_worker.py
@@ -39,12 +39,7 @@
       self.verify_suite = 'build_RootFilesystemSize'
     else:
       self.verify_suite = 'suite:%s' % (options.verify_suite_name or 'smoke')
-
-    # 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/gce_au_worker.py b/au_test_harness/gce_au_worker.py
index 74ac082..fe62db2 100644
--- a/au_test_harness/gce_au_worker.py
+++ b/au_test_harness/gce_au_worker.py
@@ -14,6 +14,7 @@
 from chromite.compute import gcloud
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_logging as logging
+from chromite.lib import path_util
 from crostestutils.au_test_harness import au_worker
 from crostestutils.au_test_harness import constants
 from crostestutils.au_test_harness import update_exception
@@ -82,7 +83,7 @@
     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):
+                  proxy_port=None, payload_signing_key=None):
     """Updates the image on a GCE instance.
 
     Unlike real_au_worker, this method always creates a new instance. Note that
@@ -117,7 +118,8 @@
     cmd = ['test_that', '--no-quickmerge', '--results_dir=%s' % test_directory,
            self.instance_ip, test]
     if self.ssh_private_key is not None:
-      cmd.append('--ssh_private_key=%s' % self.ssh_private_key)
+      cmd.append('--ssh_private_key=%s' %
+                 path_util.ToChrootPath(self.ssh_private_key))
 
     result = cros_build_lib.RunCommand(cmd, error_code_ok=True,
                                        enter_chroot=True, redirect_stdout=True,
diff --git a/au_test_harness/real_au_worker.py b/au_test_harness/real_au_worker.py
index ba7aa04..6da9ab9 100644
--- a/au_test_harness/real_au_worker.py
+++ b/au_test_harness/real_au_worker.py
@@ -6,6 +6,7 @@
 
 import constants
 from chromite.lib import cros_build_lib
+from chromite.lib import path_util
 from crostestutils.au_test_harness import au_worker
 
 
@@ -57,7 +58,8 @@
     cmd = ['test_that', '--no-quickmerge', '--results_dir=%s' % test_directory,
            self.remote, test]
     if self.ssh_private_key is not None:
-      cmd.append('--ssh_private_key=%s' % self.ssh_private_key)
+      cmd.append('--ssh_private_key=%s' %
+                 path_util.ToChrootPath(self.ssh_private_key))
 
     result = cros_build_lib.RunCommand(cmd, error_code_ok=True,
                                        enter_chroot=True, redirect_stdout=True,
diff --git a/cros_run_vm_test b/cros_run_vm_test
index 90f19f2..192883d 100755
--- a/cros_run_vm_test
+++ b/cros_run_vm_test
@@ -25,8 +25,6 @@
 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
 
@@ -85,7 +83,7 @@
 if [ -n "${FLAGS_ssh_private_key}" ]; then
   test_args+=(
     "--ssh_private_key"
-    "${FLAGS_ssh_private_key}"
+    "$(reinterpret_path_for_chroot "${FLAGS_ssh_private_key}")"
   )
 fi