Change all VM update tests to update to vm image (not non-vm image).

We have some tests now that will need to ensure that we update to the
vm image instead of the non-vm image.

I've had to make more changes to au_test_harness than originally expected
because of some ugly built-in assumptions I've cleaned up.

BUG=chromium:314970
TEST=Remote trybot.

Change-Id: I32f5a8e9a9062e79b787f630a71dafaa27a4215f
Reviewed-on: https://chromium-review.googlesource.com/174793
Reviewed-by: Chris Sosa <sosa@chromium.org>
Commit-Queue: Chris Sosa <sosa@chromium.org>
Tested-by: Chris Sosa <sosa@chromium.org>
diff --git a/au_test_harness/au_test.py b/au_test_harness/au_test.py
index 62842ab..4182f7d 100644
--- a/au_test_harness/au_test.py
+++ b/au_test_harness/au_test.py
@@ -58,7 +58,7 @@
 
   def AttemptUpdateWithFilter(self, update_filter, proxy_port=8081):
     """Update through a proxy, with a specified filter, and expect success."""
-    self.worker.PrepareBase(self.target_image_path)
+    target_image_path = self.worker.PrepareBase(self.target_image_path)
 
     # The devserver runs at port 8080 by default. We assume that here, and
     # start our proxy at a different one. We then tell our update tools to
@@ -69,7 +69,7 @@
                                           filter=update_filter)
     proxy.serve_forever_in_thread()
     try:
-      self.worker.PerformUpdate(self.target_image_path, self.target_image_path,
+      self.worker.PerformUpdate(target_image_path, target_image_path,
                                 proxy_port=proxy_port)
     finally:
       proxy.shutdown()
@@ -99,10 +99,10 @@
     self.worker.Initialize(9222)
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
-    self.worker.PrepareBase(self.base_image_path)
+    base_image_path = self.worker.PrepareBase(self.base_image_path)
 
     # Update to
-    self.worker.PerformUpdate(self.target_image_path, self.base_image_path)
+    self.worker.PerformUpdate(self.target_image_path, base_image_path)
     self.assertTrue(self.worker.VerifyImage())
 
     # Update from
@@ -118,10 +118,10 @@
     self.worker.Initialize(9223)
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
-    self.worker.PrepareBase(self.base_image_path)
+    base_image_path = self.worker.PrepareBase(self.base_image_path)
 
     # Update to
-    self.worker.PerformUpdate(self.target_image_path, self.base_image_path,
+    self.worker.PerformUpdate(self.target_image_path, base_image_path,
                               'clean')
     self.assertTrue(self.worker.VerifyImage())
 
@@ -165,10 +165,11 @@
   def testSimpleSignedUpdate(self):
     """Test that updates to itself with a signed payload."""
     self.worker.Initialize(9226)
-    self.worker.PrepareBase(self.target_image_path, signed_base=True)
+    signed_target_image_path = self.worker.PrepareBase(self.target_image_path,
+                                                       signed_base=True)
     if self.private_key:
       self.worker.PerformUpdate(self.target_image_path,
-                                self.target_image_path + '.signed',
+                                signed_target_image_path,
                                 private_key_path=self.private_key)
     else:
       cros_build_lib.Info('No key found to use for signed testing.')
@@ -180,8 +181,8 @@
     run using test_prefix option.
     """
     self.worker.Initialize(9227)
-    self.worker.PrepareBase(self.target_image_path)
-    self.worker.PerformUpdate(self.target_image_path, self.target_image_path)
+    target_image_path = self.worker.PrepareBase(self.target_image_path)
+    self.worker.PerformUpdate(target_image_path, target_image_path)
     self.assertTrue(self.worker.VerifyImage())
 
   def SimpleTestVerify(self):
diff --git a/au_test_harness/au_worker.py b/au_test_harness/au_worker.py
index 727cfcb..c1067d4 100644
--- a/au_test_harness/au_worker.py
+++ b/au_test_harness/au_worker.py
@@ -27,7 +27,6 @@
   def __init__(self, options, test_results_root):
     """Processes options for the specific-type of worker."""
     self.board = options.board
-    self._first_update = False
     self.test_results_root = test_results_root
     self.all_results_root = os.path.join(test_results_root, 'all')
     self.fail_results_root = os.path.join(test_results_root, 'failed')
@@ -60,6 +59,8 @@
     Subclasses must override this method with the correct procedure for
     preparing the test target.
 
+    Returns the path to the base image (might have changed for vm's).
+
 `   Args:
       image_path: The image that should reside on the target before the test.
       signed_base: If True, use the signed image rather than the actual image.
@@ -146,20 +147,25 @@
 
   def PrepareRealBase(self, image_path, signed_base):
     """Prepares a remote device for worker test by updating it to the image."""
+    real_image_path = image_path
     if not signed_base:
-      self.UpdateImage(image_path)
+      self.UpdateImage(real_image_path)
     else:
-      self.UpdateImage(image_path + '.signed')
+      real_image_path = real_image_path + '.signed'
+      self.UpdateImage(real_image_path)
+
+    return real_image_path
 
   def PrepareVMBase(self, image_path, signed_base):
     """Prepares a VM image for worker test."""
     # Tells the VM tests to use the Qemu image as the start point.
-    self._first_update = True
     self.vm_image_path = os.path.join(os.path.dirname(image_path),
                                       'chromiumos_qemu_image.bin')
     if signed_base:
       self.vm_image_path = self.vm_image_path + '.signed'
 
+    return self.vm_image_path
+
   def GetStatefulChangeFlag(self, stateful_change):
     """Returns the flag to pass to image_to_vm for the stateful change."""
     stateful_change_flag = ''
diff --git a/au_test_harness/real_au_worker.py b/au_test_harness/real_au_worker.py
index fb9321e..df5e698 100644
--- a/au_test_harness/real_au_worker.py
+++ b/au_test_harness/real_au_worker.py
@@ -21,7 +21,7 @@
 
   def PrepareBase(self, image_path, signed_base=False):
     """Auto-update to base image to prepare for test."""
-    self.PrepareRealBase(image_path, signed_base)
+    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):
diff --git a/au_test_harness/vm_au_worker.py b/au_test_harness/vm_au_worker.py
index 52844ca..42eb41c 100644
--- a/au_test_harness/vm_au_worker.py
+++ b/au_test_harness/vm_au_worker.py
@@ -38,7 +38,7 @@
 
   def PrepareBase(self, image_path, signed_base=False):
     """Creates an update-able VM based on base image."""
-    self.PrepareVMBase(image_path, signed_base)
+    return self.PrepareVMBase(image_path, signed_base)
 
   @staticmethod
   def _HandleFail(log_directory, fail_directory):
@@ -56,10 +56,6 @@
     """Updates VM image with image_path."""
     log_directory, fail_directory = self.GetNextResultsPath('update')
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
-    if src_image_path and self._first_update:
-      src_image_path = self.vm_image_path
-      self._first_update = False
-
     cmd = ['%s/bin/cros_run_vm_update' % constants.CROSUTILS_DIR,
            '--vm_image_path=%s' % self.vm_image_path,
            '--update_log=%s' % os.path.join(log_directory, 'update_engine.log'),
diff --git a/generate_test_payloads/cros_generate_test_payloads.py b/generate_test_payloads/cros_generate_test_payloads.py
index 1e32c76..faf782e 100755
--- a/generate_test_payloads/cros_generate_test_payloads.py
+++ b/generate_test_payloads/cros_generate_test_payloads.py
@@ -182,23 +182,22 @@
     """Generate Payload Requirements for AUTestHarness and NPlus1 Testing."""
     if self.full_suite:
       # N-1->N.
-      self._AddUpdatePayload(self.target_no_vm, self.base, for_vm=self.vm)
+      self._AddUpdatePayload(self.target, self.base, for_vm=self.vm)
 
       # N->N after N-1->N.
-      self._AddUpdatePayload(self.target_no_vm, self.target_no_vm,
-                             for_vm=self.vm)
+      self._AddUpdatePayload(self.target, self.target, for_vm=self.vm)
 
       # N->N From VM base.
-      self._AddUpdatePayload(self.target_no_vm, self.target, for_vm=self.vm)
+      self._AddUpdatePayload(self.target, self.target, for_vm=self.vm)
 
       # Need a signed payload for the signed payload test.
       if self.target_signed:
-        self._AddUpdatePayload(self.target_no_vm, self.target_signed,
+        self._AddUpdatePayload(self.target_signed, self.target_signed,
                                self.private_key, for_vm=self.vm)
 
     if self.basic_suite:
       # Update image to itself from VM base.
-      self._AddUpdatePayload(self.target_no_vm, self.target, for_vm=self.vm)
+      self._AddUpdatePayload(self.target, self.target, for_vm=self.vm)
 
     # Add deltas for m minus 1 to n and n to n.
     if self.nplus1: