src/third_party/chromite/lib/paygen/paygen_build_lib_unittest.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """Tests for paygen_build_lib."""
   7
   8 from __future__ import print_function
   9
  10 import itertools
  11 import mox
  12 import os
  13 import shutil
  14 import tempfile
  15 import unittest
  16
  17 import fixup_path
  18 fixup_path.FixupPath()
  19
  20 from chromite.cbuildbot import cbuildbot_config
  21 from chromite.cbuildbot import commands
  22
  23 from chromite.lib import cros_build_lib
  24 from chromite.lib import cros_test_lib
  25 from chromite.lib import osutils
  26 from chromite.lib import parallel
  27
  28 from chromite.lib.paygen import download_cache
  29 from chromite.lib.paygen import gslock
  30 from chromite.lib.paygen import gslib
  31 from chromite.lib.paygen import gspaths
  32 from chromite.lib.paygen import urilib
  33 from chromite.lib.paygen import paygen_build_lib
  34 from chromite.lib.paygen import paygen_payload_lib
  35 from chromite.lib.paygen import utils
  36
  37
  38 # pylint: disable-msg=F0401
  39 from site_utils.autoupdate.lib import test_control
  40 # pylint: enable-msg=F0401
  41
  42
  43 # We access a lot of protected members during testing.
  44 # pylint: disable-msg=W0212
  45
  46
  47 class PaygenBuildLibTest(mox.MoxTestBase):
  48   """Test PaygenBuildLib class."""
  49
  50   def __init__(self, test_case_names):
  51     mox.MoxTestBase.__init__(self, test_case_names)
  52     self.work_dir = '/work/foo'
  53     self.tempdir = None
  54
  55     self.prev_image = gspaths.Image(channel='foo-channel',
  56                                     board='foo-board',
  57                                     version='1.0.0',
  58                                     key='mp')
  59     self.prev2_image = gspaths.Image(channel='foo-channel',
  60                                      board='foo-board',
  61                                      version='1.1.0',
  62                                      key='mp')
  63
  64     self.foo_build = gspaths.Build(bucket='crt',
  65                                    channel='foo-channel',
  66                                    board='foo-board',
  67                                    version='1.2.3')
  68
  69     # Create an additional 'special' image like NPO that isn't NPO,
  70     # and keyed with a weird key. It should match none of the filters.
  71     self.special_image = gspaths.Image(bucket='crt',
  72                                        channel='foo-channel',
  73                                        board='foo-board',
  74                                        version='1.2.3',
  75                                        key='foo-key',
  76                                        image_channel='special-channel')
  77
  78     self.images = self._GetBuildImages(self.foo_build)
  79     (self.basic_image, self.premp_image,
  80      self.npo_image, self.premp_npo_image) = self.images
  81
  82     self.test_image = self._GetBuildTestImage(self.foo_build)
  83
  84   def _GetPaygenBuildInstance(self, skip_test_payloads=False,
  85                               disable_tests=False):
  86     """Helper method to create a standard Paygen instance."""
  87     control_dir = None if disable_tests else '/tmp/foo'
  88
  89     return paygen_build_lib._PaygenBuild(self.foo_build, self.work_dir,
  90                                          control_dir=control_dir,
  91                                          skip_test_payloads=skip_test_payloads)
  92
  93   def _GetBuildImages(self, build):
  94     """Create basic_image, npo_image, premp_image, premp_npo_image.
  95
  96     Args:
  97       build: gspaths.Build object describing the build to create fake images
  98         for.
  99     """
 100     # NPOs should have image_version incremented, but it doesn't matter for our
 101     # testing.
 102     basic_image = gspaths.Image(key='mp-v2', **build)
 103     npo_image = gspaths.Image(key='mp-v2',
 104                               image_channel='nplusone-channel',
 105                               image_version=build.version,
 106                               **build)
 107     premp_image = gspaths.Image(key='premp', **build)
 108     premp_npo_image = gspaths.Image(key='premp',
 109                                     image_channel='nplusone-channel',
 110                                     image_version=build.version,
 111                                     **build)
 112
 113     # Code in several places depends on the order.
 114     return [basic_image, premp_image, npo_image, premp_npo_image]
 115
 116   def _GetBuildTestImage(self, build):
 117     """Returns a test image object for the build.
 118
 119     Args:
 120       build: gspaths.Build object describing the build to create fake images
 121         for.
 122     """
 123     return gspaths.UnsignedImageArchive(bucket=build.bucket,
 124                                         channel=build.channel,
 125                                         board=build.board,
 126                                         version=build.version,
 127                                         milestone='R12',
 128                                         image_type='test')
 129
 130   def testGetFlagURI(self):
 131     """Validate the helper method to create flag URIs for our current build."""
 132     paygen = self._GetPaygenBuildInstance()
 133
 134     self.assertEqual(
 135         paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK),
 136         'gs://crt/foo-channel/foo-board/1.2.3/payloads/LOCK_flag')
 137     self.assertEqual(
 138         paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP),
 139         'gs://crt/foo-channel/foo-board/1.2.3/payloads/SKIP_flag')
 140     self.assertEqual(
 141         paygen._GetFlagURI(gspaths.ChromeosReleases.FINISHED),
 142         'gs://crt/foo-channel/foo-board/1.2.3/payloads/FINISHED_flag')
 143
 144   def testFilterHelpers(self):
 145     """Test _FilterForMp helper method."""
 146
 147     # All of the filter helpers should handle empty list.
 148     self.assertEqual(paygen_build_lib._FilterForMp([]), [])
 149     self.assertEqual(paygen_build_lib._FilterForPremp([]), [])
 150     self.assertEqual(paygen_build_lib._FilterForBasic([]), [])
 151     self.assertEqual(paygen_build_lib._FilterForNpo([]), [])
 152
 153     # prev_image lets us test with an 'mp' key, instead of an 'mp-v2' key.
 154     images = list(self.images) + [self.special_image, self.prev_image]
 155
 156     self.assertEqual(paygen_build_lib._FilterForMp(images),
 157                      [self.basic_image, self.npo_image, self.prev_image])
 158
 159     self.assertEqual(paygen_build_lib._FilterForPremp(images),
 160                      [self.premp_image, self.premp_npo_image])
 161
 162     self.assertEqual(paygen_build_lib._FilterForBasic(images),
 163                      [self.basic_image, self.premp_image, self.prev_image])
 164
 165     self.assertEqual(paygen_build_lib._FilterForNpo(images),
 166                      [self.npo_image, self.premp_npo_image])
 167
 168   def testValidateExpectedBuildImages(self):
 169     """Test a function that validates expected images are found on a build."""
 170     paygen = self._GetPaygenBuildInstance()
 171
 172     # Test with basic mp image only.
 173     paygen._ValidateExpectedBuildImages(self.foo_build, (self.basic_image,))
 174
 175     # Test with basic mp and mp npo images.
 176     paygen._ValidateExpectedBuildImages(self.foo_build, (self.basic_image,
 177                                                          self.npo_image,))
 178     # Test with basic mp and premp images.
 179     paygen._ValidateExpectedBuildImages(self.foo_build, (self.basic_image,
 180                                                          self.premp_image))
 181
 182     # Test with basic mp and premp images.
 183     paygen._ValidateExpectedBuildImages(self.foo_build, (self.basic_image,
 184                                                          self.premp_image,
 185                                                          self.npo_image))
 186
 187     # Test with 4 different images.
 188     paygen._ValidateExpectedBuildImages(self.foo_build, (self.basic_image,
 189                                                          self.premp_image,
 190                                                          self.npo_image,
 191                                                          self.premp_npo_image))
 192
 193     # No images isn't valid.
 194     with self.assertRaises(paygen_build_lib.ImageMissing):
 195       paygen._ValidateExpectedBuildImages(self.foo_build, [])
 196
 197     # NPO image only isn't valid.
 198     with self.assertRaises(paygen_build_lib.ImageMissing):
 199       paygen._ValidateExpectedBuildImages(self.foo_build, (self.npo_image,))
 200
 201     # NPO without matching basic isn't valid.
 202     with self.assertRaises(paygen_build_lib.ImageMissing):
 203       paygen._ValidateExpectedBuildImages(self.foo_build,
 204                                           (self.premp_image,
 205                                            self.npo_image,
 206                                            self.premp_npo_image))
 207
 208     # More than one of the same type of image should trigger BuildCorrupt
 209     with self.assertRaises(paygen_build_lib.BuildCorrupt):
 210       paygen._ValidateExpectedBuildImages(self.foo_build, (self.basic_image,
 211                                                            self.basic_image,))
 212
 213     # Unexpected images should trigger BuildCorrupt
 214     with self.assertRaises(paygen_build_lib.BuildCorrupt):
 215       paygen._ValidateExpectedBuildImages(self.foo_build,
 216                                           (self.basic_image,
 217                                            self.npo_image,
 218                                            self.special_image))
 219
 220   def _TestDiscoverArtifacts(self, list_files_uri, list_files_result,
 221                              test_func, test_args, should_succeed,
 222                              expected_result):
 223     """Test artifact discovery using mocked gsutil results."""
 224     self.mox.StubOutWithMock(urilib, 'ListFiles')
 225     urilib.ListFiles(list_files_uri).AndReturn(list_files_result)
 226     self.mox.ReplayAll()
 227
 228     if should_succeed:
 229       self.assertEqual(test_func(*test_args), expected_result)
 230     else:
 231       self.assertRaises(expected_result, test_func, *test_args)
 232
 233   def testDiscoverImages(self):
 234     """Test _DiscoverImages."""
 235     paygen = self._GetPaygenBuildInstance()
 236     uri_base = 'gs://crt/foo-channel/foo-board/1.2.3'
 237
 238     uri_basic = os.path.join(
 239         uri_base, 'chromeos_1.2.3_foo-board_recovery_foo-channel_mp-v3.bin')
 240     uri_premp = os.path.join(
 241         uri_base, 'chromeos_1.2.3_foo-board_recovery_foo-channel_premp.bin')
 242     uri_npo = os.path.join(
 243         uri_base,
 244         'chromeos_1.2.4_foo-board_recovery_nplusone-channel_mp-v3.bin')
 245     file_list_result = [uri_basic, uri_premp, uri_npo]
 246
 247     base_image_params = {'channel': 'foo-channel',
 248                          'board': 'foo-board',
 249                          'version': '1.2.3',
 250                          'bucket': 'crt'}
 251     expected_basic = gspaths.Image(key='mp-v3', uri=uri_basic,
 252                                    **base_image_params)
 253     expected_premp = gspaths.Image(key='premp', uri=uri_premp,
 254                                    **base_image_params)
 255     expected_npo = gspaths.Image(key='mp-v3', image_channel='nplusone-channel',
 256                                  image_version='1.2.4', uri=uri_npo,
 257                                  **base_image_params)
 258     expected_result = [expected_basic, expected_premp, expected_npo]
 259
 260     self._TestDiscoverArtifacts(
 261         os.path.join(uri_base, 'chromeos_*_foo-board_recovery_*_*.bin'),
 262         file_list_result,
 263         paygen._DiscoverImages,
 264         [self.foo_build],
 265         True,
 266         expected_result)
 267
 268   def testDiscoverTestImageArchives(self):
 269     """Test _DiscoverTestImageArchives (success)."""
 270     paygen = self._GetPaygenBuildInstance()
 271     uri_base = 'gs://crt/foo-channel/foo-board/1.2.3'
 272
 273     uri_test_archive = os.path.join(
 274         uri_base, 'ChromeOS-test-R12-1.2.3-foo-board.tar.xz')
 275     file_list_result = [uri_test_archive]
 276
 277     expected_test_archive = gspaths.UnsignedImageArchive(
 278         channel='foo-channel',
 279         board='foo-board',
 280         version='1.2.3',
 281         bucket='crt',
 282         uri=uri_test_archive,
 283         milestone='R12',
 284         image_type='test')
 285     expected_result = [expected_test_archive]
 286
 287     self._TestDiscoverArtifacts(
 288         os.path.join(uri_base, 'ChromeOS-test-*-1.2.3-foo-board.tar.xz'),
 289         file_list_result,
 290         paygen._DiscoverTestImageArchives,
 291         [self.foo_build],
 292         True,
 293         expected_result)
 294
 295   def testDiscoverTestImageArchivesMultipleResults(self):
 296     """Test _DiscoverTestImageArchives (fails due to multiple results)."""
 297     paygen = self._GetPaygenBuildInstance()
 298     uri_base = 'gs://crt/foo-channel/foo-board/1.2.3'
 299
 300     uri_test_archive1 = os.path.join(
 301         uri_base, 'ChromeOS-test-R12-1.2.3-foo-board.tar.xz')
 302     uri_test_archive2 = os.path.join(
 303         uri_base, 'ChromeOS-test-R13-1.2.3-foo-board.tar.xz')
 304     file_list_result = [uri_test_archive1, uri_test_archive2]
 305
 306     self._TestDiscoverArtifacts(
 307         os.path.join(uri_base, 'ChromeOS-test-*-1.2.3-foo-board.tar.xz'),
 308         file_list_result,
 309         paygen._DiscoverTestImageArchives,
 310         [self.foo_build],
 311         False,
 312         paygen_build_lib.BuildCorrupt)
 313
 314   def testDiscoverTestImageArchivesMissing(self):
 315     """Test _DiscoverTestImageArchives (fails due to missing images)."""
 316     paygen = self._GetPaygenBuildInstance()
 317     uri_base = 'gs://crt/foo-channel/foo-board/1.2.3'
 318
 319     self._TestDiscoverArtifacts(
 320         os.path.join(uri_base, 'ChromeOS-test-*-1.2.3-foo-board.tar.xz'),
 321         [],
 322         paygen._DiscoverTestImageArchives,
 323         [self.foo_build],
 324         False,
 325         paygen_build_lib.ImageMissing)
 326
 327   @unittest.skipIf(not paygen_build_lib.config, 'Internal crostools required.')
 328   def testDiscoverFsiBuilds(self):
 329     """Using test release.conf values, test _DiscoverFsiBuilds."""
 330
 331     test_config = """
 332 [valid-board]
 333 fsi_images: 2913.331.0,2465.105.0
 334
 335 [no-fsi-board]
 336 """
 337     paygen_build_lib.config.LoadTestConfig(test_config)
 338
 339     # Test a board with FSI values on stable-channel.
 340     paygen = paygen_build_lib._PaygenBuild(
 341         gspaths.Build(channel='stable-channel', board='valid-board',
 342                       version='1.2.3'),
 343         self.work_dir)
 344
 345     self.assertEqual(
 346         sorted(paygen._DiscoverFsiBuilds()),
 347         [gspaths.Build(board='valid-board',
 348                        channel='stable-channel',
 349                        version='2465.105.0'),
 350          gspaths.Build(board='valid-board',
 351                        channel='stable-channel',
 352                        version='2913.331.0')])
 353
 354     # Test a board without FSI values on stable-channel.
 355     paygen = paygen_build_lib._PaygenBuild(
 356         gspaths.Build(channel='stable-channel', board='no-fsi-board',
 357                       version='1.2.3'),
 358         self.work_dir)
 359
 360     self.assertEqual(paygen._DiscoverFsiBuilds(), [])
 361
 362     # Test a board with FSI values on non-stable-channel.
 363     paygen = paygen_build_lib._PaygenBuild(
 364         gspaths.Build(channel='beta-channel', board='valid-board',
 365                       version='1.2.3'),
 366         self.work_dir)
 367
 368     self.assertEqual(paygen._DiscoverFsiBuilds(), [])
 369
 370     paygen_build_lib.config.LoadGlobalConfig()
 371
 372   @unittest.skipIf(not paygen_build_lib.query, 'Internal crostools required.')
 373   def testDiscoverNmoBuild(self):
 374     """Test _DiscoverNmoBuild (N minus One)."""
 375     paygen = self._GetPaygenBuildInstance()
 376
 377     self.mox.StubOutWithMock(paygen_build_lib.query, 'FindLatestPublished')
 378
 379     # Set up the test replay script.
 380     paygen_build_lib.query.FindLatestPublished(
 381         'foo-channel', 'foo-board').AndReturn('1.0.0')
 382
 383     paygen_build_lib.query.FindLatestPublished(
 384         'foo-channel', 'foo-board').AndReturn(None)
 385
 386     # Run the test verification.
 387     self.mox.ReplayAll()
 388
 389     self.assertEqual(paygen._DiscoverNmoBuild(),
 390                      [gspaths.Build(bucket='crt',
 391                                     channel='foo-channel',
 392                                     board='foo-board',
 393                                     version='1.0.0')])
 394
 395     self.assertEqual(paygen._DiscoverNmoBuild(), [])
 396
 397   def testDiscoverRequiredFullPayloads(self):
 398     """Test _DiscoverRequiredFullPayloads."""
 399     paygen = self._GetPaygenBuildInstance()
 400
 401     self.assertEqual(paygen._DiscoverRequiredFullPayloads([]), [])
 402
 403     self.assertItemsEqual(
 404         paygen._DiscoverRequiredFullPayloads(self.images + [self.test_image]),
 405         [gspaths.Payload(tgt_image=self.basic_image),
 406          gspaths.Payload(tgt_image=self.npo_image),
 407          gspaths.Payload(tgt_image=self.premp_image),
 408          gspaths.Payload(tgt_image=self.premp_npo_image),
 409          gspaths.Payload(tgt_image=self.test_image)])
 410
 411   def testDiscoverRequiredNpoDeltas(self):
 412     """Test _DiscoverRequiredNpoDeltas."""
 413     paygen = self._GetPaygenBuildInstance()
 414
 415     self.assertEqual(paygen._DiscoverRequiredNpoDeltas([]), [])
 416
 417     self.assertEqual(paygen._DiscoverRequiredNpoDeltas([self.basic_image]), [])
 418
 419     self.assertEqual(paygen._DiscoverRequiredNpoDeltas([self.npo_image]), [])
 420
 421     expected = [gspaths.Payload(tgt_image=self.npo_image,
 422                                 src_image=self.basic_image)]
 423     self.assertEqual(paygen._DiscoverRequiredNpoDeltas([self.basic_image,
 424                                                         self.npo_image]),
 425                      expected)
 426
 427     self.assertEqual(paygen._DiscoverRequiredNpoDeltas([self.npo_image,
 428                                                         self.basic_image]),
 429                      expected)
 430
 431     self.assertEqual(paygen._DiscoverRequiredNpoDeltas([self.premp_image,
 432                                                         self.premp_npo_image]),
 433                      [gspaths.Payload(tgt_image=self.premp_npo_image,
 434                                       src_image=self.premp_image)])
 435
 436   def testDiscoverRequiredTestNpoDeltas(self):
 437     """Test _DiscoverRequiredTestNpoDeltas."""
 438     paygen = self._GetPaygenBuildInstance()
 439
 440     self.assertEqual(
 441         paygen._DiscoverRequiredTestNpoDeltas([]), [])
 442     self.assertItemsEqual(
 443         paygen._DiscoverRequiredTestNpoDeltas([self.test_image]),
 444         [gspaths.Payload(tgt_image=self.test_image,
 445                          src_image=self.test_image)])
 446
 447   def testDiscoverRequiredFromPreviousDeltas(self):
 448     """Test _DiscoverRequiredFromPreviousDeltas."""
 449     paygen = self._GetPaygenBuildInstance()
 450
 451     images = [self.basic_image]
 452     prevs = [self.prev_image, self.prev2_image]
 453
 454     # Empty lists.
 455     results = paygen._DiscoverRequiredFromPreviousDeltas([], [])
 456     expected = []
 457     self.assertEqual(results, expected)
 458
 459     # Empty previous list.
 460     results = paygen._DiscoverRequiredFromPreviousDeltas(images, [])
 461     expected = []
 462     self.assertEqual(results, expected)
 463
 464     # Empty target list.
 465     results = paygen._DiscoverRequiredFromPreviousDeltas([], prevs)
 466     expected = []
 467     self.assertEqual(results, expected)
 468
 469     # Basic list.
 470     results = paygen._DiscoverRequiredFromPreviousDeltas(images, prevs)
 471     expected = [gspaths.Payload(tgt_image=self.basic_image,
 472                                 src_image=self.prev_image),
 473                 gspaths.Payload(tgt_image=self.basic_image,
 474                                 src_image=self.prev2_image)]
 475     self.assertEqual(results, expected)
 476
 477     # Inverted order (should return nothing).
 478     results = paygen._DiscoverRequiredFromPreviousDeltas(
 479         [self.prev_image], images)
 480     expected = []
 481     self.assertEqual(results, expected)
 482
 483   def testDiscoverRequiredPayloadsIncompleteBuild(self):
 484     """Test _DiscoverRequiredPayloads."""
 485
 486     paygen = self._GetPaygenBuildInstance()
 487
 488     self.mox.StubOutWithMock(paygen, '_DiscoverImages')
 489     self.mox.StubOutWithMock(paygen, '_DiscoverNmoBuild')
 490     self.mox.StubOutWithMock(paygen, '_DiscoverFsiBuilds')
 491
 492     paygen.BUILD_DISCOVER_RETRY_SLEEP = 0
 493
 494     # Check that we retry 3 times.
 495     paygen._DiscoverImages(paygen._build).AndRaise(
 496         paygen_build_lib.ImageMissing())
 497     paygen._DiscoverImages(paygen._build).AndRaise(
 498         paygen_build_lib.ImageMissing())
 499     paygen._DiscoverImages(paygen._build).AndRaise(
 500         paygen_build_lib.ImageMissing())
 501     paygen._DiscoverImages(paygen._build).AndRaise(
 502         paygen_build_lib.ImageMissing())
 503
 504     # Run the test verification.
 505     self.mox.ReplayAll()
 506
 507     with self.assertRaises(paygen_build_lib.BuildNotReady):
 508       paygen._DiscoverRequiredPayloads()
 509
 510   def testDiscoverRequiredPayloads(self):
 511     """Test _DiscoverRequiredPayloads."""
 512
 513     paygen = self._GetPaygenBuildInstance()
 514
 515     output_uri = 'gs://foo'
 516
 517     self.mox.StubOutWithMock(paygen, '_DiscoverImages')
 518     self.mox.StubOutWithMock(paygen, '_DiscoverTestImageArchives')
 519     self.mox.StubOutWithMock(paygen, '_DiscoverNmoBuild')
 520     self.mox.StubOutWithMock(paygen, '_DiscoverFsiBuilds')
 521     self.mox.StubOutWithMock(paygen_payload_lib, 'DefaultPayloadUri')
 522
 523     nmo_build = gspaths.Build(bucket='crt',
 524                               channel='foo-channel',
 525                               board='foo-board',
 526                               version='1.2.2')
 527     fsi1_build = gspaths.Build(bucket='crt',
 528                                channel='foo-channel',
 529                                board='foo-board',
 530                                version='1.0.0')
 531     fsi2_build = gspaths.Build(bucket='crt',
 532                                channel='foo-channel',
 533                                board='foo-board',
 534                                version='1.1.0')
 535
 536     nmo_images = self._GetBuildImages(nmo_build)
 537     nmo_test_image = self._GetBuildTestImage(nmo_build)
 538     fsi1_images = self._GetBuildImages(fsi1_build)
 539     fsi1_test_image = self._GetBuildTestImage(fsi1_build)
 540     fsi2_images = self._GetBuildImages(fsi2_build)
 541     fsi2_test_image = self._GetBuildTestImage(fsi2_build)
 542
 543     paygen._DiscoverImages(paygen._build).AndReturn(self.images)
 544     paygen._DiscoverTestImageArchives(paygen._build).AndReturn(
 545         [self.test_image])
 546     paygen._DiscoverNmoBuild().AndReturn([nmo_build])
 547     paygen._DiscoverFsiBuilds().AndReturn([fsi1_build, fsi2_build])
 548     paygen._DiscoverImages(nmo_build).AndReturn(nmo_images)
 549     paygen._DiscoverTestImageArchives(nmo_build).AndReturn([nmo_test_image])
 550     paygen._DiscoverImages(fsi1_build).AndReturn(fsi1_images)
 551     paygen._DiscoverTestImageArchives(fsi1_build).AndReturn([fsi1_test_image])
 552     paygen._DiscoverImages(fsi2_build).AndReturn(fsi2_images)
 553     paygen._DiscoverTestImageArchives(fsi2_build).AndReturn([fsi2_test_image])
 554
 555     # Simplify the output URIs, so it's easy to check them below.
 556     paygen_payload_lib.DefaultPayloadUri(
 557         mox.IsA(gspaths.Payload), None).MultipleTimes().AndReturn(output_uri)
 558
 559     # Run the test verification.
 560     self.mox.ReplayAll()
 561
 562     results = paygen._DiscoverRequiredPayloads()
 563
 564     expected = [gspaths.Payload(tgt_image=self.basic_image, uri=output_uri),
 565                 gspaths.Payload(tgt_image=self.npo_image, uri=output_uri),
 566                 gspaths.Payload(tgt_image=self.premp_image, uri=output_uri),
 567                 gspaths.Payload(tgt_image=self.premp_npo_image, uri=output_uri),
 568                 # NPO Deltas
 569                 gspaths.Payload(tgt_image=self.npo_image,
 570                                 src_image=self.basic_image,
 571                                 uri=output_uri),
 572                 gspaths.Payload(tgt_image=self.premp_npo_image,
 573                                 src_image=self.premp_image,
 574                                 uri=output_uri),
 575                 # NMO Delta
 576                 gspaths.Payload(tgt_image=self.basic_image,
 577                                 src_image=nmo_images[0],
 578                                 uri=output_uri),
 579                 gspaths.Payload(tgt_image=self.premp_image,
 580                                 src_image=nmo_images[1],
 581                                 uri=output_uri),
 582                 # FSI Deltas
 583                 gspaths.Payload(tgt_image=self.basic_image,
 584                                 src_image=fsi1_images[0],
 585                                 uri=output_uri),
 586                 gspaths.Payload(tgt_image=self.premp_image,
 587                                 src_image=fsi1_images[1],
 588                                 uri=output_uri),
 589                 gspaths.Payload(tgt_image=self.basic_image,
 590                                 src_image=fsi2_images[0],
 591                                 uri=output_uri),
 592                 gspaths.Payload(tgt_image=self.premp_image,
 593                                 src_image=fsi2_images[1],
 594                                 uri=output_uri),
 595
 596                 # Test full payload.
 597                 gspaths.Payload(tgt_image=self.test_image,
 598                                 uri=output_uri),
 599
 600                 # Test NPO delta.
 601                 gspaths.Payload(tgt_image=self.test_image,
 602                                 src_image=self.test_image,
 603                                 uri=output_uri),
 604
 605                 # Test NMO delta.
 606                 gspaths.Payload(tgt_image=self.test_image,
 607                                 src_image=nmo_test_image,
 608                                 uri=output_uri),
 609
 610                 # Test FSI deltas.
 611                 gspaths.Payload(tgt_image=self.test_image,
 612                                 src_image=fsi1_test_image,
 613                                 uri=output_uri),
 614                 gspaths.Payload(tgt_image=self.test_image,
 615                                 src_image=fsi2_test_image,
 616                                 uri=output_uri),
 617                ]
 618     expected = zip(expected, itertools.repeat(False))
 619
 620     self.assertItemsEqual(sorted(results), sorted(expected))
 621
 622   def testDiscoverRequiredPayloadsPreviousSkipped(self):
 623     """Test _DiscoverRequiredPayload.
 624
 625     Ensures that no test delta payload is generated if generation of a
 626     signed delta from the same build was skipped.
 627     """
 628
 629     paygen = self._GetPaygenBuildInstance()
 630
 631     output_uri = 'gs://foo'
 632
 633     self.mox.StubOutWithMock(paygen, '_DiscoverImages')
 634     self.mox.StubOutWithMock(paygen, '_DiscoverTestImageArchives')
 635     self.mox.StubOutWithMock(paygen, '_DiscoverNmoBuild')
 636     self.mox.StubOutWithMock(paygen, '_DiscoverFsiBuilds')
 637     self.mox.StubOutWithMock(paygen_payload_lib, 'DefaultPayloadUri')
 638
 639     nmo_build = gspaths.Build(bucket='crt',
 640                               channel='foo-channel',
 641                               board='foo-board',
 642                               version='1.2.2')
 643     fsi1_build = gspaths.Build(bucket='crt',
 644                                channel='foo-channel',
 645                                board='foo-board',
 646                                version='1.0.0')
 647     fsi2_build = gspaths.Build(bucket='crt',
 648                                channel='foo-channel',
 649                                board='foo-board',
 650                                version='1.1.0')
 651
 652     fsi1_images = self._GetBuildImages(fsi1_build)
 653     fsi1_test_image = self._GetBuildTestImage(fsi1_build)
 654     fsi2_images = self._GetBuildImages(fsi2_build)
 655     fsi2_test_image = self._GetBuildTestImage(fsi2_build)
 656
 657     paygen._DiscoverImages(paygen._build).AndReturn(self.images)
 658     paygen._DiscoverTestImageArchives(paygen._build).AndReturn(
 659         [self.test_image])
 660     paygen._DiscoverNmoBuild().AndReturn([nmo_build])
 661     paygen._DiscoverFsiBuilds().AndReturn([fsi1_build, fsi2_build])
 662     paygen._DiscoverImages(nmo_build).AndRaise(
 663         paygen_build_lib.ImageMissing('nmo build is missing some image'))
 664     # _DiscoverTestImageArchives(nmo_build) should NOT be called.
 665     paygen._DiscoverImages(fsi1_build).AndReturn(fsi1_images)
 666     paygen._DiscoverTestImageArchives(fsi1_build).AndReturn([fsi1_test_image])
 667     paygen._DiscoverImages(fsi2_build).AndReturn(fsi2_images)
 668     paygen._DiscoverTestImageArchives(fsi2_build).AndReturn([fsi2_test_image])
 669
 670     # Simplify the output URIs, so it's easy to check them below.
 671     paygen_payload_lib.DefaultPayloadUri(
 672         mox.IsA(gspaths.Payload), None).MultipleTimes().AndReturn(output_uri)
 673
 674     # Run the test verification.
 675     self.mox.ReplayAll()
 676
 677     results = paygen._DiscoverRequiredPayloads()
 678
 679     # IMPORTANT: we intentionally omit the NMO payload from the expected list
 680     # of payloads as it is a duplicate of one of the FSIs.
 681     expected = [gspaths.Payload(tgt_image=self.basic_image, uri=output_uri),
 682                 gspaths.Payload(tgt_image=self.npo_image, uri=output_uri),
 683                 gspaths.Payload(tgt_image=self.premp_image, uri=output_uri),
 684                 gspaths.Payload(tgt_image=self.premp_npo_image, uri=output_uri),
 685                 # NPO Deltas
 686                 gspaths.Payload(tgt_image=self.npo_image,
 687                                 src_image=self.basic_image,
 688                                 uri=output_uri),
 689                 gspaths.Payload(tgt_image=self.premp_npo_image,
 690                                 src_image=self.premp_image,
 691                                 uri=output_uri),
 692                 # FSI Deltas
 693                 gspaths.Payload(tgt_image=self.basic_image,
 694                                 src_image=fsi1_images[0],
 695                                 uri=output_uri),
 696                 gspaths.Payload(tgt_image=self.premp_image,
 697                                 src_image=fsi1_images[1],
 698                                 uri=output_uri),
 699                 gspaths.Payload(tgt_image=self.basic_image,
 700                                 src_image=fsi2_images[0],
 701                                 uri=output_uri),
 702                 gspaths.Payload(tgt_image=self.premp_image,
 703                                 src_image=fsi2_images[1],
 704                                 uri=output_uri),
 705
 706                 # Test full payload.
 707                 gspaths.Payload(tgt_image=self.test_image,
 708                                 uri=output_uri),
 709
 710                 # Test NPO delta.
 711                 gspaths.Payload(tgt_image=self.test_image,
 712                                 src_image=self.test_image,
 713                                 uri=output_uri),
 714
 715                 # Test FSI deltas.
 716                 gspaths.Payload(tgt_image=self.test_image,
 717                                 src_image=fsi1_test_image,
 718                                 uri=output_uri),
 719                 gspaths.Payload(tgt_image=self.test_image,
 720                                 src_image=fsi2_test_image,
 721                                 uri=output_uri),
 722                ]
 723     expected = zip(expected, itertools.repeat(False))
 724
 725     self.assertItemsEqual(sorted(results), sorted(expected))
 726
 727   def testDiscoverRequiredPayloadsNmoIsAlsoFsi(self):
 728     """Test _DiscoverRequiredPayloads."""
 729
 730     paygen = self._GetPaygenBuildInstance()
 731
 732     output_uri = 'gs://foo'
 733
 734     self.mox.StubOutWithMock(paygen, '_DiscoverImages')
 735     self.mox.StubOutWithMock(paygen, '_DiscoverTestImageArchives')
 736     self.mox.StubOutWithMock(paygen, '_DiscoverNmoBuild')
 737     self.mox.StubOutWithMock(paygen, '_DiscoverFsiBuilds')
 738     self.mox.StubOutWithMock(paygen_payload_lib, 'DefaultPayloadUri')
 739
 740     nmo_build = gspaths.Build(bucket='crt',
 741                               channel='foo-channel',
 742                               board='foo-board',
 743                               version='1.2.2')
 744     fsi1_build = gspaths.Build(bucket='crt',
 745                                channel='foo-channel',
 746                                board='foo-board',
 747                                version='1.0.0')
 748     fsi2_build = gspaths.Build(bucket='crt',
 749                                channel='foo-channel',
 750                                board='foo-board',
 751                                version='1.2.2')
 752
 753     nmo_images = self._GetBuildImages(nmo_build)
 754     nmo_test_image = self._GetBuildTestImage(nmo_build)
 755     fsi1_images = self._GetBuildImages(fsi1_build)
 756     fsi1_test_image = self._GetBuildTestImage(fsi1_build)
 757
 758     paygen._DiscoverImages(paygen._build).AndReturn(self.images)
 759     paygen._DiscoverTestImageArchives(paygen._build).AndReturn(
 760         [self.test_image])
 761     paygen._DiscoverNmoBuild().AndReturn([nmo_build])
 762     paygen._DiscoverFsiBuilds().AndReturn([fsi1_build, fsi2_build])
 763     paygen._DiscoverImages(nmo_build).AndReturn(nmo_images)
 764     paygen._DiscoverTestImageArchives(nmo_build).AndReturn([nmo_test_image])
 765     paygen._DiscoverImages(fsi1_build).AndReturn(fsi1_images)
 766     paygen._DiscoverTestImageArchives(fsi1_build).AndReturn([fsi1_test_image])
 767
 768     # Simplify the output URIs, so it's easy to check them below.
 769     paygen_payload_lib.DefaultPayloadUri(
 770         mox.IsA(gspaths.Payload), None).MultipleTimes().AndReturn(output_uri)
 771
 772     # Run the test verification.
 773     self.mox.ReplayAll()
 774
 775     results = paygen._DiscoverRequiredPayloads()
 776
 777     expected = [gspaths.Payload(tgt_image=self.basic_image, uri=output_uri),
 778                 gspaths.Payload(tgt_image=self.npo_image, uri=output_uri),
 779                 gspaths.Payload(tgt_image=self.premp_image, uri=output_uri),
 780                 gspaths.Payload(tgt_image=self.premp_npo_image, uri=output_uri),
 781                 # NPO Deltas
 782                 gspaths.Payload(tgt_image=self.npo_image,
 783                                 src_image=self.basic_image,
 784                                 uri=output_uri),
 785                 gspaths.Payload(tgt_image=self.premp_npo_image,
 786                                 src_image=self.premp_image,
 787                                 uri=output_uri),
 788                 # NMO Delta
 789                 gspaths.Payload(tgt_image=self.basic_image,
 790                                 src_image=nmo_images[0],
 791                                 uri=output_uri),
 792                 gspaths.Payload(tgt_image=self.premp_image,
 793                                 src_image=nmo_images[1],
 794                                 uri=output_uri),
 795                 # FSI Deltas
 796                 gspaths.Payload(tgt_image=self.basic_image,
 797                                 src_image=fsi1_images[0],
 798                                 uri=output_uri),
 799                 gspaths.Payload(tgt_image=self.premp_image,
 800                                 src_image=fsi1_images[1],
 801                                 uri=output_uri),
 802
 803                 # Test full payload.
 804                 gspaths.Payload(tgt_image=self.test_image,
 805                                 uri=output_uri),
 806
 807                 # Test NPO delta.
 808                 gspaths.Payload(tgt_image=self.test_image,
 809                                 src_image=self.test_image,
 810                                 uri=output_uri),
 811
 812                 # Test NMO delta.
 813                 gspaths.Payload(tgt_image=self.test_image,
 814                                 src_image=nmo_test_image,
 815                                 uri=output_uri),
 816
 817                 # Test FSI deltas.
 818                 gspaths.Payload(tgt_image=self.test_image,
 819                                 src_image=fsi1_test_image,
 820                                 uri=output_uri),
 821                ]
 822     expected = zip(expected, itertools.repeat(False))
 823
 824     self.assertItemsEqual(sorted(results), sorted(expected))
 825
 826   def testFindFullTestPayloads(self):
 827     paygen = self._GetPaygenBuildInstance()
 828
 829     self.mox.StubOutWithMock(urilib, 'ListFiles')
 830
 831     urilib.ListFiles(
 832         'gs://crt/foo-channel/foo-board/find_full_version/payloads/'
 833         'chromeos_find_full_version_foo-board_foo-channel_full_test.bin-*'
 834     ).AndReturn(['foo', 'foo.json', 'foo.log', 'bar'])
 835
 836     # Run the test verification.
 837     self.mox.ReplayAll()
 838
 839     # Call once and use mocked look up. Make sure we filter properly.
 840     self.assertEqual(paygen._FindFullTestPayloads('find_full_version'),
 841                      ['foo', 'bar'])
 842
 843     # Call a second time to verify we get cached results (no lookup).
 844     self.assertEqual(paygen._FindFullTestPayloads('find_full_version'),
 845                      ['foo', 'bar'])
 846
 847   @osutils.TempDirDecorator
 848   def DoGeneratePayloadsTest(self, run_parallel, test_dry_run):
 849     """Test paygen_build_lib._GeneratePayloads."""
 850     paygen = paygen_build_lib._PaygenBuild(self.foo_build, self.tempdir,
 851                                            dry_run=test_dry_run,
 852                                            run_parallel=run_parallel)
 853
 854     basic_payload = gspaths.Payload(tgt_image=self.npo_image,
 855                                     src_image=self.basic_image)
 856     premp_payload = gspaths.Payload(tgt_image=self.premp_npo_image,
 857                                     src_image=self.premp_image)
 858
 859     self.mox.StubOutWithMock(parallel, 'RunTasksInProcessPool')
 860     self.mox.StubOutWithMock(paygen_build_lib, '_GenerateSinglePayload')
 861
 862     expected_payload_args = [
 863         (basic_payload, mox.IsA(str), True, test_dry_run),
 864         (premp_payload, mox.IsA(str), True, test_dry_run)
 865         ]
 866
 867     if run_parallel:
 868       parallel.RunTasksInProcessPool(paygen_build_lib._GenerateSinglePayload,
 869                                      expected_payload_args)
 870     else:
 871       paygen_build_lib._GenerateSinglePayload(basic_payload, mox.IsA(str),
 872                                               True, test_dry_run)
 873
 874       paygen_build_lib._GenerateSinglePayload(premp_payload, mox.IsA(str),
 875                                               True, test_dry_run)
 876
 877     # Run the test verification.
 878     self.mox.ReplayAll()
 879
 880     paygen._GeneratePayloads((basic_payload, premp_payload), lock=None)
 881
 882     self.mox.UnsetStubs()
 883
 884   def testGeneratePayloads(self):
 885     """Test paygen_build_lib._GeneratePayloads, no dry_run."""
 886
 887     # Test every combination of the boolean arguments.
 888     for run_parallel in (True, False):
 889       for test_dry_run in (True, False):
 890         self.DoGeneratePayloadsTest(run_parallel, test_dry_run)
 891
 892   @osutils.TempDirDecorator
 893   def testGeneratePayloadInProcess(self):
 894     """Make sure the _GenerateSinglePayload calls into paygen_payload_lib."""
 895
 896     basic_payload = gspaths.Payload(tgt_image=self.npo_image,
 897                                     src_image=self.basic_image)
 898
 899     self.mox.StubOutWithMock(paygen_payload_lib, 'CreateAndUploadPayload')
 900
 901     # Verify that we actually generate the payload.
 902     paygen_payload_lib.CreateAndUploadPayload(
 903         basic_payload,
 904         mox.IsA(download_cache.DownloadCache),
 905         work_dir=self.tempdir,
 906         sign=False,
 907         dry_run=True)
 908
 909     # Run the test verification.
 910     self.mox.ReplayAll()
 911
 912     paygen_build_lib._GenerateSinglePayload(basic_payload, self.tempdir,
 913                                             sign=False, dry_run=True)
 914
 915   def testCleanupBuild(self):
 916     """Test _PaygenBuild._CleanupBuild."""
 917     paygen = self._GetPaygenBuildInstance()
 918
 919     self.mox.StubOutWithMock(gslib, 'Remove')
 920     gslib.Remove('gs://crt/foo-channel/foo-board/1.2.3/payloads/signing',
 921                  recurse=True, ignore_no_match=True)
 922     self.mox.ReplayAll()
 923
 924     paygen._CleanupBuild()
 925
 926   def _CreatePayloadsSetup(self, skip_test_payloads=False, disable_tests=False):
 927     """Helper method for related CreatePayloads tests."""
 928     paygen = self._GetPaygenBuildInstance(skip_test_payloads=skip_test_payloads,
 929                                           disable_tests=disable_tests)
 930
 931     self.mox.StubOutWithMock(gslock, 'Lock')
 932     self.mox.StubOutWithMock(gslib, 'CreateWithContents')
 933     self.mox.StubOutWithMock(gslib, 'Exists')
 934     self.mox.StubOutWithMock(gslib, 'Remove')
 935     self.mox.StubOutWithMock(paygen, '_DiscoverRequiredPayloads')
 936     self.mox.StubOutWithMock(paygen, '_MapToArchive')
 937     self.mox.StubOutWithMock(paygen, '_GeneratePayloads')
 938     self.mox.StubOutWithMock(paygen, '_AutotestPayloads')
 939     self.mox.StubOutWithMock(paygen, '_CreatePayloadTests')
 940     self.mox.StubOutWithMock(paygen, '_CleanupBuild')
 941
 942     return paygen
 943
 944   def testCreatePayloadsLockedBuild(self):
 945     """Test paygen_build_lib._GeneratePayloads if the build is locked."""
 946     paygen = self._CreatePayloadsSetup()
 947     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
 948
 949     # Pylint is confused because of helper stubbing things out.
 950     # pylint: disable-msg=E1101
 951     gslock.Lock(lock_uri, dry_run=False).AndRaise(gslock.LockNotAcquired())
 952
 953     # Run the test verification.
 954     self.mox.ReplayAll()
 955
 956     with self.assertRaises(paygen_build_lib.BuildLocked):
 957       paygen.CreatePayloads()
 958
 959   def testCreatePayloadsSkipBuild(self):
 960     """Test paygen_build_lib._GeneratePayloads if the build marked skip."""
 961     paygen = self._CreatePayloadsSetup()
 962     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
 963     skip_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP)
 964
 965     lock = self.mox.CreateMockAnything()
 966
 967     gslock.Lock(lock_uri, dry_run=False).AndReturn(lock)
 968     lock.__enter__().AndReturn(lock)
 969     gslib.Exists(skip_uri).AndReturn(True)
 970     lock.__exit__(
 971         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
 972
 973     # Run the test verification.
 974     self.mox.ReplayAll()
 975
 976     with self.assertRaises(paygen_build_lib.BuildSkip):
 977       paygen.CreatePayloads()
 978
 979   def testCreatePayloadsFinishedBuild(self):
 980     """Test paygen_build_lib._GeneratePayloads if the build marked finished."""
 981     paygen = self._CreatePayloadsSetup()
 982
 983     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
 984     skip_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP)
 985     finished_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
 986
 987     lock = self.mox.CreateMockAnything()
 988
 989     gslock.Lock(lock_uri, dry_run=False).AndReturn(lock)
 990     lock.__enter__().AndReturn(lock)
 991     gslib.Exists(skip_uri).AndReturn(False)
 992     gslib.Exists(finished_uri).AndReturn(True)
 993     lock.__exit__(
 994         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
 995
 996     # Run the test verification.
 997     self.mox.ReplayAll()
 998
 999     with self.assertRaises(paygen_build_lib.BuildFinished):
1000       paygen.CreatePayloads()
1001
1002   def testCreatePayloadsBuildNotReady(self):
1003     """Test paygen_build_lib._GeneratePayloads if not all images are there."""
1004     paygen = self._CreatePayloadsSetup()
1005
1006     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
1007     skip_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP)
1008     finished_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
1009
1010     lock = self.mox.CreateMockAnything()
1011
1012     gslock.Lock(lock_uri, dry_run=False).AndReturn(lock)
1013     lock.__enter__().AndReturn(lock)
1014     gslib.Exists(skip_uri).AndReturn(False)
1015     gslib.Exists(finished_uri).AndReturn(False)
1016     # This method is being mocked out.
1017     # pylint: disable=E1101
1018     paygen._DiscoverRequiredPayloads(
1019         ).AndRaise(paygen_build_lib.BuildNotReady())
1020     lock.__exit__(
1021         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
1022
1023     # Run the test verification.
1024     self.mox.ReplayAll()
1025
1026     with self.assertRaises(paygen_build_lib.BuildNotReady):
1027       paygen.CreatePayloads()
1028
1029   def testCreatePayloadsCreateFailed(self):
1030     """Test paygen_build_lib._GeneratePayloads if payload generation failed."""
1031     paygen = self._CreatePayloadsSetup()
1032
1033     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
1034     skip_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP)
1035     finished_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
1036
1037     lock = self.mox.CreateMockAnything()
1038     payload = 'foo'
1039     payload_list = [payload]
1040     payload_skip_list = [(payload, False)]
1041     mock_exception = Exception()
1042
1043     gslock.Lock(lock_uri, dry_run=False).AndReturn(lock)
1044     lock.__enter__().AndReturn(lock)
1045     gslib.Exists(skip_uri).AndReturn(False)
1046     gslib.Exists(finished_uri).AndReturn(False)
1047     paygen._DiscoverRequiredPayloads(
1048         ).AndReturn(payload_skip_list)
1049     self.mox.StubOutWithMock(paygen_payload_lib, 'FindExistingPayloads')
1050     paygen_payload_lib.FindExistingPayloads(payload).AndReturn([])
1051     paygen._GeneratePayloads(payload_list, lock).AndRaise(mock_exception)
1052     lock.__exit__(
1053         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
1054
1055     # Run the test verification.
1056     self.mox.ReplayAll()
1057
1058     with self.assertRaises(Exception):
1059       paygen.CreatePayloads()
1060
1061   def testCreatePayloadsSuccess(self):
1062     """Test paygen_build_lib._GeneratePayloads success."""
1063     paygen = self._CreatePayloadsSetup()
1064
1065     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
1066     skip_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP)
1067     finished_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
1068
1069     lock = self.mox.CreateMockAnything()
1070     payload = 'foo'
1071     payload_list = [payload]
1072     payload_skip_list = [(payload, False)]
1073
1074     gslock.Lock(lock_uri, dry_run=False).AndReturn(lock)
1075     lock.__enter__().AndReturn(lock)
1076     gslib.Exists(skip_uri).AndReturn(False)
1077     gslib.Exists(finished_uri).AndReturn(False)
1078     paygen._DiscoverRequiredPayloads(
1079         ).AndReturn(payload_skip_list)
1080     self.mox.StubOutWithMock(paygen_payload_lib, 'FindExistingPayloads')
1081     paygen_payload_lib.FindExistingPayloads(payload).AndReturn([])
1082     paygen._GeneratePayloads(payload_list, lock)
1083     paygen._MapToArchive('foo-board', '1.2.3').AndReturn(
1084         ('archive_board', 'archive_build', 'archive_build_uri'))
1085     paygen._CreatePayloadTests(['foo']).AndReturn(['Test Payloads'])
1086     paygen._AutotestPayloads(['Test Payloads'])
1087
1088     paygen._CleanupBuild()
1089     gslib.CreateWithContents(finished_uri, mox.IgnoreArg())
1090     lock.__exit__(
1091         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
1092
1093     # Run the test verification.
1094     self.mox.ReplayAll()
1095
1096     paygen.CreatePayloads()
1097
1098   def testCreatePayloadsAlreadyExists(self):
1099     """Test paygen_build_lib._GeneratePayloads success."""
1100     paygen = self._CreatePayloadsSetup()
1101
1102     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
1103     skip_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP)
1104     finished_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
1105
1106     lock = self.mox.CreateMockAnything()
1107     self.mox.StubOutWithMock(paygen_payload_lib, 'FindExistingPayloads')
1108     self.mox.StubOutWithMock(paygen_payload_lib, 'SetPayloadUri')
1109     payload_existing = 'foo'
1110     payload_new = 'bar'
1111     payload_list = [(payload_existing, False), (payload_new, False)]
1112
1113     gslock.Lock(lock_uri, dry_run=False).AndReturn(lock)
1114     lock.__enter__().AndReturn(lock)
1115     gslib.Exists(skip_uri).AndReturn(False)
1116     gslib.Exists(finished_uri).AndReturn(False)
1117     paygen._DiscoverRequiredPayloads(
1118         ).AndReturn(payload_list)
1119     paygen_payload_lib.FindExistingPayloads(payload_existing).AndReturn(
1120         [payload_existing])
1121     paygen_payload_lib.FindExistingPayloads(payload_new).AndReturn([])
1122     paygen_payload_lib.SetPayloadUri(payload_existing, payload_existing)
1123     paygen._GeneratePayloads([payload_new], lock)
1124     paygen._MapToArchive('foo-board', '1.2.3').AndReturn(
1125         ('archive_board', 'archive_build', 'archive_build_uri'))
1126     paygen._CreatePayloadTests(['foo', 'bar']).AndReturn(['Test Payloads'])
1127     paygen._AutotestPayloads(['Test Payloads'])
1128     gslib.CreateWithContents(finished_uri, mox.IgnoreArg())
1129     paygen._CleanupBuild()
1130     lock.__exit__(
1131         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
1132
1133     # Run the test verification.
1134     self.mox.ReplayAll()
1135
1136     paygen.CreatePayloads()
1137
1138   def testCreatePayloadsSkipTests(self):
1139     """Test paygen_build_lib._GeneratePayloads success."""
1140     paygen = self._CreatePayloadsSetup(skip_test_payloads=True,
1141                                        disable_tests=True)
1142
1143     lock_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.LOCK)
1144     skip_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.SKIP)
1145     finished_uri = paygen._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
1146
1147     lock = self.mox.CreateMockAnything()
1148     payload = 'foo'
1149     payload_list = [payload]
1150     payload_skip_list = [(payload, False)]
1151
1152     gslock.Lock(lock_uri, dry_run=False).AndReturn(lock)
1153     lock.__enter__().AndReturn(lock)
1154     gslib.Exists(skip_uri).AndReturn(False)
1155     gslib.Exists(finished_uri).AndReturn(False)
1156     paygen._DiscoverRequiredPayloads(
1157         ).AndReturn(payload_skip_list)
1158     self.mox.StubOutWithMock(paygen_payload_lib, 'FindExistingPayloads')
1159     paygen_payload_lib.FindExistingPayloads(payload).AndReturn([])
1160     paygen._GeneratePayloads(payload_list, lock)
1161     paygen._CleanupBuild()
1162     gslib.CreateWithContents(finished_uri, mox.IgnoreArg())
1163     lock.__exit__(
1164         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
1165
1166     # Run the test verification.
1167     self.mox.ReplayAll()
1168
1169     paygen.CreatePayloads()
1170
1171   @osutils.TempDirDecorator
1172   def testFindControlFileDir(self):
1173     """Test that we find control files in the proper directory."""
1174     # Test default dir in /tmp.
1175     result = paygen_build_lib._FindControlFileDir(None)
1176     self.assertTrue(os.path.isdir(result))
1177     self.assertTrue(result.startswith('/tmp/paygen_build-control_files'))
1178     shutil.rmtree(result)
1179
1180     # Test in specified dir.
1181     result = paygen_build_lib._FindControlFileDir(self.tempdir)
1182     self.assertTrue(os.path.isdir(result))
1183     self.assertTrue(result.startswith(
1184         os.path.join(self.tempdir, 'paygen_build-control_files')))
1185
1186   @unittest.skipIf(not paygen_build_lib.config,
1187                    'Internal crostools repository needed.')
1188   def testEmitControlFile(self):
1189     """Test that we emit control files correctly."""
1190     payload = gspaths.Payload(tgt_image=self.npo_image,
1191                               src_image=self.basic_image)
1192
1193     suite_name = 'paygen_foo'
1194     control_dir = tempfile.mkdtemp(prefix='control_dir-')
1195     paygen = paygen_build_lib._PaygenBuild(
1196         self.foo_build, self.tempdir, control_dir=control_dir)
1197     with tempfile.NamedTemporaryFile(prefix='control_file-', delete=False) as f:
1198       control_file_name = f.name
1199       f.write("""
1200 AUTHOR = "Chromium OS"
1201 NAME = "autoupdate_EndToEndTest"
1202 TIME = "MEDIUM"
1203 TEST_CATEGORY = "Functional"
1204 TEST_CLASS = "platform"
1205 TEST_TYPE = "server"
1206 DOC = "Faux doc"
1207
1208 """)
1209
1210     self.mox.StubOutWithMock(cbuildbot_config, 'FindFullConfigsForBoard')
1211     cbuildbot_config.FindFullConfigsForBoard().AndReturn(
1212         ([{'boards': ['foo_board']}], []))
1213
1214     self.mox.StubOutWithMock(urilib, 'ListFiles')
1215     urilib.ListFiles(
1216         gspaths.ChromeosReleases.PayloadUri(
1217             self.basic_image.channel, self.basic_image.board,
1218             self.basic_image.version,
1219             '*', bucket=self.basic_image.bucket)).AndReturn(
1220                 ['gs://foo/bar.tar.bz2'])
1221     urilib.ListFiles(
1222         gspaths.ChromeosImageArchive.BuildUri(
1223             'foo_board', '*', self.basic_image.version)).AndReturn(
1224                 ['gs://foo-archive/src-build'])
1225
1226     self.mox.StubOutWithMock(test_control, 'get_control_file_name')
1227     test_control.get_control_file_name().AndReturn(control_file_name)
1228
1229     self.mox.ReplayAll()
1230
1231     payload_test = paygen_build_lib._PaygenBuild.PayloadTest(payload)
1232     paygen._EmitControlFile(payload_test, suite_name, control_dir)
1233
1234     shutil.rmtree(control_dir)
1235     os.remove(control_file_name)
1236
1237   def testAutotestPayloads(self):
1238     """Test the process of scheduling HWLab tests."""
1239     control_dir = '/tmp/control_dir'
1240     paygen = paygen_build_lib._PaygenBuild(
1241         self.foo_build, self.tempdir, control_dir=control_dir)
1242     control_dump_dir = os.path.join(control_dir, paygen.CONTROL_FILE_SUBDIR)
1243     payloads = ['foo', 'bar']
1244     test_channel = self.foo_build.channel.rpartition('-')[0]
1245     suite_name = paygen.PAYGEN_AU_SUITE_TEMPLATE % test_channel
1246     tarball_name = paygen.CONTROL_TARBALL_TEMPLATE % test_channel
1247     tarball_path = os.path.join(control_dir, tarball_name)
1248     test_archive_build = '%s-release/R99-%s' % (self.foo_build.board,
1249                                                 self.foo_build.version)
1250     test_archive_build_uri = ('gs://chromeos-image-archive/%s' %
1251                               test_archive_build)
1252     test_upload_path = os.path.join(test_archive_build_uri, tarball_name)
1253
1254     self.mox.StubOutWithMock(os, 'makedirs')
1255     os.makedirs(os.path.join(control_dir, paygen.CONTROL_FILE_SUBDIR))
1256
1257     self.mox.StubOutWithMock(paygen, '_EmitControlFile')
1258     paygen._EmitControlFile('foo', suite_name, control_dump_dir)
1259     paygen._EmitControlFile('bar', suite_name, control_dump_dir)
1260
1261     self.mox.StubOutWithMock(cros_build_lib, 'CreateTarball')
1262     cros_build_lib.CreateTarball(
1263         tarball_path, control_dir,
1264         compression=cros_build_lib.COMP_BZIP2,
1265         inputs=[paygen.CONTROL_FILE_SUBDIR]).AndReturn(
1266             cros_build_lib.CommandResult(returncode=0))
1267
1268     # Setup preliminary values needed for running autotests.
1269     paygen._archive_board = self.foo_build.board
1270     paygen._archive_build = test_archive_build
1271     paygen._archive_build_uri = test_archive_build_uri
1272
1273     self.mox.StubOutWithMock(gslib, 'Copy')
1274     gslib.Copy(tarball_path, test_upload_path, acl='public-read')
1275
1276     # TODO(garnold) remove the dryrun argument.
1277     self.mox.StubOutWithMock(utils, 'RunCommand')
1278
1279     timeout_mins = cbuildbot_config.HWTestConfig.DEFAULT_HW_TEST_TIMEOUT / 60
1280     expected_command = [
1281         mox.StrContains('site_utils/run_suite.py'),
1282         '--board', 'foo-board',
1283         '--build', 'foo-board-release/R99-1.2.3',
1284         '--suite_name', 'paygen_au_foo',
1285         '--file_bugs', 'True',
1286         '--pool', 'bvt',
1287         '--retry', 'True',
1288         '--timeout_mins', str(timeout_mins),
1289         '--no_wait', 'False']
1290
1291     utils.RunCommand(expected_command, error_ok=True, redirect_stdout=True,
1292                      redirect_stderr=True, return_result=True).AndReturn(
1293                          utils.CommandResult(returncode=0))
1294
1295     self.mox.ReplayAll()
1296
1297     paygen._AutotestPayloads(payloads)
1298
1299   def testScheduleAutotestTestsNormal(self):
1300     """Test scheduling autotest tests with run_suite.py."""
1301     paygen = paygen_build_lib._PaygenBuild(
1302         self.foo_build, self.tempdir)
1303
1304     self.mox.StubOutWithMock(paygen_build_lib.commands,
1305                              'RunHWTestSuite')
1306     self.mox.StubOutWithMock(paygen_build_lib.utils,
1307                              'RunCommand')
1308
1309     timeout_mins = cbuildbot_config.HWTestConfig.DEFAULT_HW_TEST_TIMEOUT / 60
1310     expected_command = [
1311         mox.StrContains('site_utils/run_suite.py'),
1312         '--board', 'foo-board',
1313         '--build', 'foo-board-release/R99-1.2.3',
1314         '--suite_name', 'paygen_au_foo',
1315         '--file_bugs', 'True',
1316         '--pool', 'bvt',
1317         '--retry', 'True',
1318         '--timeout_mins', str(timeout_mins),
1319         '--no_wait', 'False']
1320     utils.RunCommand(expected_command, error_ok=True, redirect_stdout=True,
1321                      redirect_stderr=True, return_result=True).AndReturn(
1322                          utils.CommandResult(returncode=0))
1323
1324     self.mox.ReplayAll()
1325
1326     # Setup preliminary values needed for scheduling autotests.
1327     paygen._archive_board = 'foo-board'
1328     paygen._archive_build = 'foo-board-release/R99-1.2.3'
1329
1330     paygen._ScheduleAutotestTests('paygen_au_foo')
1331
1332   def testScheduleAutotestTestsBuilderEnvironment(self):
1333     """Test scheduling autotest tests with build autotest proxy."""
1334     paygen = paygen_build_lib._PaygenBuild(
1335         self.foo_build, self.tempdir, run_on_builder=True)
1336
1337     self.mox.StubOutWithMock(paygen_build_lib.commands,
1338                              'RunHWTestSuite')
1339     self.mox.StubOutWithMock(paygen_build_lib.utils,
1340                              'RunCommand')
1341
1342     timeout_mins = cbuildbot_config.HWTestConfig.DEFAULT_HW_TEST_TIMEOUT / 60
1343     paygen_build_lib.commands.RunHWTestSuite(
1344         board='foo-board', build='foo-board-release/R99-1.2.3', file_bugs=True,
1345         pool='bvt', suite='paygen_au_foo', timeout_mins=timeout_mins,
1346         retry=True, wait_for_results=True, debug=False)
1347
1348     self.mox.ReplayAll()
1349
1350     # Setup preliminary values needed for scheduling autotests.
1351     paygen._archive_board = 'foo-board'
1352     paygen._archive_build = 'foo-board-release/R99-1.2.3'
1353
1354     paygen._ScheduleAutotestTests('paygen_au_foo')
1355
1356   def testScheduleAutotestTestsBuilderEnvironmentWarn(self):
1357     """Test scheduling autotest tests with build autotest proxy."""
1358     paygen = paygen_build_lib._PaygenBuild(
1359         self.foo_build, self.tempdir, run_on_builder=True)
1360
1361     self.mox.StubOutWithMock(paygen_build_lib.commands,
1362                              'RunHWTestSuite')
1363     self.mox.StubOutWithMock(paygen_build_lib.utils,
1364                              'RunCommand')
1365
1366     timeout_mins = cbuildbot_config.HWTestConfig.DEFAULT_HW_TEST_TIMEOUT / 60
1367     paygen_build_lib.commands.RunHWTestSuite(
1368         board='foo-board', build='foo-board-release/R99-1.2.3', file_bugs=True,
1369         pool='bvt', suite='paygen_au_foo', timeout_mins=timeout_mins,
1370         retry=True, wait_for_results=True, debug=False).AndRaise(
1371           commands.TestWarning('** Suite passed with a warning code **'))
1372
1373     self.mox.ReplayAll()
1374
1375     # Setup preliminary values needed for scheduling autotests.
1376     paygen._archive_board = 'foo-board'
1377     paygen._archive_build = 'foo-board-release/R99-1.2.3'
1378
1379     paygen._ScheduleAutotestTests('paygen_au_foo')
1380
1381   def testMapToArchive(self):
1382     """Test that mapping to images archive names/locations works."""
1383     self.mox.StubOutWithMock(cbuildbot_config, 'FindFullConfigsForBoard')
1384     cbuildbot_config.FindFullConfigsForBoard().MultipleTimes().AndReturn(
1385         ([{'boards': ['foo_board', 'bar_board', 'bar-board']}], []))
1386
1387     self.mox.StubOutWithMock(urilib, 'ListFiles')
1388     urilib.ListFiles(
1389         gspaths.ChromeosImageArchive.BuildUri(
1390             'foo_board', '*', '1.2.3')).AndReturn(
1391                 ['gs://foo-archive/foo_board/R11-1.2.3/somefile'])
1392
1393     self.mox.ReplayAll()
1394
1395     # Case 1: mapping successful.
1396     self.assertEqual(
1397         paygen_build_lib._PaygenBuild._MapToArchive('foo-board', '1.2.3'),
1398         ('foo_board', 'foo_board/R11-1.2.3',
1399          'gs://foo-archive/foo_board/R11-1.2.3'))
1400
1401     # Case 2: failure, too many build board names found.
1402     with self.assertRaises(paygen_build_lib.ArchiveError):
1403       paygen_build_lib._PaygenBuild._MapToArchive('bar-board', '1.2.3')
1404
1405     # Case 3: failure, build board name not found.
1406     with self.assertRaises(paygen_build_lib.ArchiveError):
1407       paygen_build_lib._PaygenBuild._MapToArchive('baz-board', '1.2.3')
1408
1409   def testValidateBoardConfig(self):
1410     """Test ValidateBoardConfig."""
1411
1412     # If we are running on an external builder, we can't see the config.
1413     # Without the config, we can't validate.
1414     if not paygen_build_lib.config:
1415       return
1416
1417     # Test a known board works.
1418     paygen_build_lib.ValidateBoardConfig('x86-mario')
1419
1420     # Test an unknown board doesn't.
1421     self.assertRaises(paygen_build_lib.BoardNotConfigured,
1422                       paygen_build_lib.ValidateBoardConfig, 'goofy-board')
1423
1424
1425 if __name__ == '__main__':
1426   cros_test_lib.main()