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