paygen_build_lib: Use stateful.tgz from chromeos-releases.

Attempt to use stateful.tgz from gs://chromeos-releases, if it's
present. Other CLs will populate this file for new builds, and manual
copies will populate it for older builds that we need to support.

If the new location for stateful isn't populated, find it in the old
chromeos-image-archive location, and copy it over. After a while,
all stateful.tgz files we'll need will have been copied forward and
we won't need the fallback mechanism any more.

I plan to merge this CL to active release branches so that we'll copy
needed files forward as quickly as possible.

Also, added unittests to ensure CreateTestPayload works after CL broke
in production the first time.

BUG=chromium:523122
TEST=Unittests

Change-Id: I55615284195e9fc5c9cb872daa89e5406d114745
Reviewed-on: https://chromium-review.googlesource.com/300571
Reviewed-by: Don Garrett <dgarrett@chromium.org>
Commit-Queue: Don Garrett <dgarrett@chromium.org>
Tested-by: Don Garrett <dgarrett@chromium.org>
diff --git a/lib/paygen/paygen_build_lib.py b/lib/paygen/paygen_build_lib.py
index 265d7d7..c892144 100644
--- a/lib/paygen/paygen_build_lib.py
+++ b/lib/paygen/paygen_build_lib.py
@@ -294,17 +294,34 @@
   class PayloadTest(object):
     """A payload test definition.
 
+    You must either use a delta payload, or specify both the src_channel and
+    src_version.
+
     Attrs:
       payload: A gspaths.Payload object describing the payload to be tested.
-      src_channel: The channel of the image to test updating from.
-      src_version: The version of the image to test updating from.
+
+      src_channel: The channel of the image to test updating from. Required
+                   if the payload is a full payload, required to be None if
+                   it's a delta.
+      src_version: The version of the image to test updating from. Required
+                   if the payload is a full payload, required to be None if
+                   it's a delta.
         for a delta payload, as it already encodes the source version.
     """
-
     def __init__(self, payload, src_channel=None, src_version=None):
       self.payload = payload
-      self.src_channel = src_channel
-      self.src_version = src_version
+
+      assert bool(src_channel) == bool(src_version), (
+          'src_channel(%s), src_version(%s) must both be set, or not set' %
+          (src_channel, src_version))
+
+      assert bool(src_channel and src_version) ^ bool(payload.src_image), (
+          'src_channel(%s), src_version(%s) required for full, not allowed'
+          ' for deltas. src_image: %s ' %
+          (src_channel, src_version, payload.src_image))
+
+      self.src_channel = src_channel or payload.src_image.channel
+      self.src_version = src_version or payload.src_image.version
 
     def __str__(self):
       return ('<test for %s%s>' %
@@ -315,6 +332,11 @@
     def __repr__(self):
       return str(self)
 
+    def __eq__(self, other):
+      return (self.payload == other.payload and
+              self.src_channel == other.src_channel and
+              self.src_version == other.src_version)
+
   def __init__(self, build, work_dir, site_config,
                dry_run=False, ignore_finished=False,
                skip_full_payloads=False, skip_delta_payloads=False,
@@ -896,14 +918,14 @@
     multiple tests in a single run.
 
     Args:
-      channel: Channel to look in for payload. If None, use build's channel.
+      channel: Channel to look in for payload.
       version: A build version whose payloads to look for.
 
     Returns:
       A (possibly empty) list of payload URIs.
     """
-    if channel is None:
-      channel = self._build.channel
+    assert channel
+    assert version
 
     if (channel, version) in self._version_to_full_test_payloads:
       # Serve from cache, if possible.
@@ -930,13 +952,6 @@
     payload = payload_test.payload
     src_version = payload_test.src_version
     src_channel = payload_test.src_channel
-    if not src_version:
-      if not payload.src_image:
-        raise PayloadTestError(
-            'no source version provided for testing full payload %s' %
-            payload)
-
-      src_version = payload.src_image.version
 
     # Discover the full test payload that corresponds to the source version.
     src_payload_uri_list = self._FindFullTestPayloads(src_channel, src_version)
@@ -957,13 +972,22 @@
     src_payload_uri = src_payload_uri_list[0]
     logging.info('Source full test payload found at %s', src_payload_uri)
 
-    # Find the chromeos_image_archive location of the source build.
-    try:
-      _, _, source_archive_uri = self._MapToArchive(
-          payload.tgt_image.board, src_version)
-    except ArchiveError as e:
-      raise PayloadTestError(
-          'error mapping source build to images archive: %s' % e)
+    release_archive_uri = gspaths.ChromeosReleases.BuildUri(
+        src_channel, self._build.board, src_version)
+
+    # TODO(dgarrett): Remove if block after finishing crbug.com/523122
+    if not urilib.Exists(os.path.join(release_archive_uri, 'stateful.tgz')):
+      logging.warning('Falling back to chromeos-image-archive: %s', payload)
+      try:
+        _, _, source_archive_uri = self._MapToArchive(
+            payload.tgt_image.board, src_version)
+      except ArchiveError as e:
+        raise PayloadTestError(
+            'error mapping source build to images archive: %s' % e)
+      stateful_archive_uri = os.path.join(source_archive_uri, 'stateful.tgz')
+      logging.info('Copying stateful.tgz from %s -> %s',
+                   stateful_archive_uri, release_archive_uri)
+      urilib.Copy(stateful_archive_uri, release_archive_uri)
 
     test = test_params.TestConfig(
         self._archive_board,
@@ -975,13 +999,14 @@
         src_payload_uri,
         payload.uri,
         suite_name=suite_name,
-        source_archive_uri=source_archive_uri)
+        source_archive_uri=release_archive_uri)
 
     with open(test_control.get_control_file_name()) as f:
       control_code = f.read()
     control_file = test_control.dump_autotest_control_file(
         test, None, control_code, control_dump_dir)
     logging.info('Control file emitted at %s', control_file)
+    return control_file
 
   def _ScheduleAutotestTests(self, suite_name):
     """Run the appropriate command to schedule the Autotests we have prepped.
@@ -1175,11 +1200,14 @@
               'update test.', self._previous_version, payload)
         else:
           payload_tests.append(self.PayloadTest(
-              payload, src_version=self._previous_version))
+              payload, src_channel=self._build.channel,
+              src_version=self._previous_version))
 
         # Create a full update test from the current version to itself.
         payload_tests.append(self.PayloadTest(
-            payload, src_channel=None, src_version=self._build.version))
+            payload,
+            src_channel=self._build.channel,
+            src_version=self._build.version))
 
         # Create a full update test from oldest viable FSI.
         payload_tests += self._CreateFsiPayloadTests(
diff --git a/lib/paygen/paygen_build_lib_unittest.py b/lib/paygen/paygen_build_lib_unittest.py
index d11abe7..20ef679 100644
--- a/lib/paygen/paygen_build_lib_unittest.py
+++ b/lib/paygen/paygen_build_lib_unittest.py
@@ -21,6 +21,7 @@
 
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_test_lib
+from chromite.lib import osutils
 from chromite.lib import parallel
 
 from chromite.lib.paygen import download_cache
@@ -1182,6 +1183,52 @@
 
     paygen.CreatePayloads()
 
+  def setupCreatePayloadTests(self):
+    paygen = self._GetPaygenBuildInstance()
+
+    self.mox.StubOutWithMock(paygen, '_DiscoverAllFsiBuilds')
+    self.mox.StubOutWithMock(paygen, '_FindFullTestPayloads')
+
+    return paygen
+
+  def testCreatePayloadTestsEmpty(self):
+
+    payloads = []
+    paygen = self.setupCreatePayloadTests()
+
+    # Run the test verification.
+    self.mox.ReplayAll()
+
+    expected = paygen._CreatePayloadTests(payloads)
+    self.assertEqual(expected, [])
+
+  def testCreatePayloadTestsPopulated(self):
+
+    payloads = [
+        gspaths.Payload(tgt_image=self.test_image),
+        gspaths.Payload(tgt_image=self.prev_image, src_image=self.test_image)
+    ]
+    paygen = self.setupCreatePayloadTests()
+
+    # We search for FSIs once for each full payload.
+    paygen._DiscoverAllFsiBuilds().AndReturn(['0.9.9', '1.0.0'])
+    paygen._FindFullTestPayloads('stable-channel', '0.9.9').AndReturn(False)
+    paygen._FindFullTestPayloads('stable-channel', '1.0.0').AndReturn(True)
+
+    # Run the test verification.
+    self.mox.ReplayAll()
+
+    self.maxDiff = None
+
+    expected = paygen._CreatePayloadTests(payloads)
+    self.assertEqual(expected, [
+        paygen.PayloadTest(
+            payloads[0], src_channel='foo-channel', src_version='1.2.3'),
+        paygen.PayloadTest(
+            payloads[0], src_channel='stable-channel', src_version='1.0.0'),
+    ])
+
+
   def testFindControlFileDir(self):
     """Test that we find control files in the proper directory."""
     # Test default dir in /tmp.
@@ -1229,16 +1276,18 @@
 """)
 
     self.mox.StubOutWithMock(urilib, 'ListFiles')
+    self.mox.StubOutWithMock(urilib, 'Exists')
+
+    urilib.Exists(
+        'gs://chromeos-releases/foo-channel/foo-board/1.2.3/stateful.tgz'
+        ).AndReturn(True)
+
     urilib.ListFiles(
         gspaths.ChromeosReleases.PayloadUri(
             self.basic_image.channel, self.basic_image.board,
             self.basic_image.version,
             '*', bucket=self.basic_image.bucket)).AndReturn(
                 ['gs://foo/bar.tar.bz2'])
-    urilib.ListFiles(
-        gspaths.ChromeosImageArchive.BuildUri(
-            'foo_board', '*', self.basic_image.version)).AndReturn(
-                ['gs://foo-archive/src-build'])
 
     self.mox.StubOutWithMock(
         paygen_build_lib.test_control, 'get_control_file_name')
@@ -1248,7 +1297,29 @@
     self.mox.ReplayAll()
 
     payload_test = paygen_build_lib._PaygenBuild.PayloadTest(payload)
-    paygen._EmitControlFile(payload_test, suite_name, control_dir)
+    cf = paygen._EmitControlFile(payload_test, suite_name, control_dir)
+
+    control_contents = osutils.ReadFile(cf)
+
+    self.assertEqual(control_contents, '''name = 'paygen_foo'
+image_type = 'test'
+update_type = 'delta'
+source_release = '1.2.3'
+target_release = '1.2.3'
+source_image_uri = 'gs://foo/bar.tar.bz2'
+target_payload_uri = 'None'
+SUITE = 'paygen_foo'
+source_archive_uri = 'gs://chromeos-releases/foo-channel/foo-board/1.2.3'
+
+AUTHOR = "Chromium OS"
+NAME = "autoupdate_EndToEndTest_paygen_foo_delta_1.2.3"
+TIME = "MEDIUM"
+TEST_CATEGORY = "Functional"
+TEST_CLASS = "platform"
+TEST_TYPE = "server"
+DOC = "Faux doc"
+
+''')
 
     shutil.rmtree(control_dir)
     os.remove(control_file_name)