Paygen: Split up Paygen into Build + Test stages.

Paygen was one giant stage.  Don had done some of the work to split out
the signer portion into its own stage.  This commit aims to split out
the testing portion into its own stage.

I call it, PaygenTest.  This stage essentially just schedules payload
tests and waits for the results.  The test payloads and artifacts are
created as part of the PaygenBuild stage and the PaygenTest stage is not
created until the artifacts are uploaded to GS.

The PaygenBuild and PaygenTest stages will be child stages of the Paygen
stage.

This will also be helpful in distinguishing failures with the actual
payload generation and those of the results from testing (or trying to
in some cases).

BUG=chromium:588282
BRANCH=None
TEST=python release_stages_unittest
TEST=python paygen_build_lib_unittest --network
TEST=payload tests tryjob

Change-Id: I4eafbd019af2c880adacaefa66bcdf393a60a1f7
Signed-off-by: Aseda Aboagye <aaboagye@google.com>
Reviewed-on: https://chromium-review.googlesource.com/388922
Commit-Ready: Aseda Aboagye <aaboagye@chromium.org>
Tested-by: Aseda Aboagye <aaboagye@chromium.org>
Reviewed-by: Don Garrett <dgarrett@chromium.org>
diff --git a/cbuildbot/stages/release_stages.py b/cbuildbot/stages/release_stages.py
index 8449774..2d97e5f 100644
--- a/cbuildbot/stages/release_stages.py
+++ b/cbuildbot/stages/release_stages.py
@@ -8,6 +8,7 @@
 
 import json
 import os
+import socket
 
 from chromite.cbuildbot import commands
 from chromite.lib import failures_lib
@@ -19,6 +20,8 @@
 from chromite.lib import osutils
 from chromite.lib import parallel
 from chromite.lib import timeout_util
+from chromite.lib.paygen import dryrun_lib
+from chromite.lib.paygen import gslib
 from chromite.lib.paygen import gspaths
 from chromite.lib.paygen import paygen_build_lib
 
@@ -326,7 +329,7 @@
 
   def _HandleStageException(self, exc_info):
     """Override and don't set status to FAIL but FORGIVEN instead."""
-    exc_type, exc_value, _exc_tb = exc_info
+    exc_type, _exc_value, _exc_tb = exc_info
 
     # If Paygen fails to find anything needed in release.conf, treat it
     # as a warning. This is common during new board bring up.
@@ -337,16 +340,6 @@
     # outright. Let SigningStage decide if this should kill the build.
     if issubclass(exc_type, SignerFailure):
       return self._HandleExceptionAsWarning(exc_info)
-
-    # If the exception is a TestLabFailure that means we couldn't schedule the
-    # test. We don't fail the build for that. We do the CompoundFailure dance,
-    # because that's how we'll get failures from background processes returned
-    # to us.
-    if (issubclass(exc_type, failures_lib.TestLabFailure) or
-        (issubclass(exc_type, failures_lib.CompoundFailure) and
-         exc_value.MatchesFailureType(failures_lib.TestLabFailure))):
-      return self._HandleExceptionAsWarning(exc_info)
-
     return super(PaygenStage, self)._HandleStageException(exc_info)
 
   def WaitUntilReady(self):
@@ -387,14 +380,14 @@
           'Golden Eye (%s) has no entry for board %s. Get a TPM to fix.' %
           (paygen_build_lib.BOARDS_URI, board))
 
+    # Default to False, set to True if it's a canary type build
+    skip_duts_check = False
+    if config_lib.IsCanaryType(self._run.config.build_type):
+      skip_duts_check = True
+
     with parallel.BackgroundTaskRunner(self._RunPaygenInProcess) as per_channel:
       logging.info("Using channels: %s", self.channels)
 
-      # Default to False, set to True if it's a canary type build
-      skip_duts_check = False
-      if config_lib.IsCanaryType(self._run.config.build_type):
-        skip_duts_check = True
-
       # If we have an explicit list of channels, use it.
       for channel in self.channels:
         per_channel.put((channel, board, version, self._run.debug,
@@ -404,41 +397,72 @@
 
   def _RunPaygenInProcess(self, channel, board, version, debug,
                           disable_tests, skip_delta_payloads,
-                          skip_duts_check=False):
-    """Helper for PaygenStage that invokes payload generation.
+                          skip_duts_check):
+    """Runs the PaygenBuild and PaygenTest stage (if applicable)"""
+    PaygenBuildStage(self._run, board, channel, version, debug, disable_tests,
+                     skip_delta_payloads, skip_duts_check).Run()
 
-    This method is intended to be safe to invoke inside a process.
+class PaygenBuildStage(generic_stages.BoardSpecificBuilderStage):
+  """Stage that generates payloads and uploads to Google Storage."""
+  def __init__(self, builder_run, board, channel, version, debug,
+               skip_testing, skip_delta_payloads, skip_duts_check, **kwargs):
+    """Init that accepts the channels argument, if present.
 
     Args:
-      channel: Channel of payloads to generate ('stable', 'beta', etc)
+      builder_run: See builder_run on ArchiveStage
       board: Board of payloads to generate ('x86-mario', 'x86-alex-he', etc)
+      channel: Channel of payloads to generate ('stable', 'beta', etc)
       version: Version of payloads to generate.
       debug: Flag telling if this is a real run, or a test run.
-      disable_tests: Do not generate test artifacts are run payload tests.
+      skip_testing: Do not generate test artifacts or run payload tests.
       skip_delta_payloads: Skip generating delta payloads.
-      skip_duts_check: Do not check minimum available DUTs
+      skip_duts_check: Do not check minimum available DUTs before tests.
+    """
+    super(PaygenBuildStage, self).__init__(builder_run, board, **kwargs)
+    self._run = builder_run
+    self.board = board
+    self.channel = channel
+    self.version = version
+    self.debug = debug
+    self.skip_testing = skip_testing
+    self.skip_delta_payloads = skip_delta_payloads
+    self.skip_duts_check = skip_duts_check
+
+  def PerformStage(self):
+    """Invoke payload generation. If testing is enabled, schedule tests.
+
+    This method is intended to be safe to invoke inside a process.
     """
     # Convert to release tools naming for channels.
-    if not channel.endswith('-channel'):
-      channel += '-channel'
+    if not self.channel.endswith('-channel'):
+      self.channel += '-channel'
 
     with osutils.TempDir(sudo_rm=True) as tempdir:
       # Create the definition of the build to generate payloads for.
-      build = gspaths.Build(channel=channel,
-                            board=board,
-                            version=version)
+      build = gspaths.Build(channel=self.channel,
+                            board=self.board,
+                            version=self.version)
 
       try:
         # Generate the payloads.
-        self._PrintLoudly('Starting %s, %s, %s' % (channel, version, board))
-        paygen_build_lib.CreatePayloads(build,
-                                        work_dir=tempdir,
-                                        site_config=self._run.site_config,
-                                        dry_run=debug,
-                                        run_parallel=True,
-                                        skip_delta_payloads=skip_delta_payloads,
-                                        disable_tests=disable_tests,
-                                        skip_duts_check=skip_duts_check)
+        self._PrintLoudly('Starting %s, %s, %s' % (self.channel, self.version,
+                                                   self.board))
+        metadata = paygen_build_lib.CreatePayloads(
+            build,
+            work_dir=tempdir,
+            site_config=self._run.site_config,
+            dry_run=self.debug,
+            run_parallel=True,
+            skip_delta_payloads=self.skip_delta_payloads,
+            disable_tests=self.skip_testing,
+            skip_duts_check=self.skip_duts_check)
+        suite_name, archive_board, archive_build, finished_uri = metadata
+
+        # Now, schedule the payload tests if desired.
+        if not self.skip_testing:
+          PaygenTestStage(suite_name, archive_board, archive_build,
+                          finished_uri, self.skip_duts_check, self._run).Run()
+
       except (paygen_build_lib.BuildFinished,
               paygen_build_lib.BuildLocked) as e:
         # These errors are normal if it's possible that another builder is, or
@@ -447,4 +471,43 @@
         #
         # This means the build was finished by the other process, or is already
         # being processed (so the build is locked).
-        logging.info('Paygen skipped because: %s', e)
+        logging.info('PaygenBuild for %s skipped because: %s', self.channel, e)
+
+
+class PaygenTestStage(generic_stages.BoardSpecificBuilderStage):
+  """Stage that schedules the payload tests."""
+  def __init__(self, suite_name, board, build, finished_uri, skip_duts_check,
+               builder_run, **kwargs):
+    self.suite_name = suite_name
+    self.board = board
+    self.build = build
+    self.finished_uri = finished_uri
+    self.skip_duts_check = skip_duts_check
+    super(PaygenTestStage, self).__init__(builder_run, board, **kwargs)
+    self._drm = dryrun_lib.DryRunMgr(self._run.debug)
+
+  def PerformStage(self):
+    """Schedule the tests to run."""
+    # Schedule the tests to run and wait for the results.
+    paygen_build_lib.ScheduleAutotestTests(self.suite_name, self.board,
+                                           self.build, self.skip_duts_check,
+                                           self._run.debug)
+
+    # Mark the build as finished since the payloads were generated, uploaded,
+    # and tested by this point.
+    self._drm(gslib.CreateWithContents, self.finished_uri, socket.gethostname())
+
+  def _HandleStageException(self, exc_info):
+    """Override and don't set status to FAIL but FORGIVEN instead."""
+    exc_type, exc_value, _exc_tb = exc_info
+
+    # If the exception is a TestLabFailure that means we couldn't schedule the
+    # test. We don't fail the build for that. We do the CompoundFailure dance,
+    # because that's how we'll get failures from background processes returned
+    # to us.
+    if (issubclass(exc_type, failures_lib.TestLabFailure) or
+        (issubclass(exc_type, failures_lib.CompoundFailure) and
+         exc_value.MatchesFailureType(failures_lib.TestLabFailure))):
+      return self._HandleExceptionAsWarning(exc_info)
+
+    return super(PaygenTestStage, self)._HandleStageException(exc_info)
diff --git a/cbuildbot/stages/release_stages_unittest.py b/cbuildbot/stages/release_stages_unittest.py
index 7edcf3c..3c66b3a 100644
--- a/cbuildbot/stages/release_stages_unittest.py
+++ b/cbuildbot/stages/release_stages_unittest.py
@@ -291,7 +291,7 @@
 
 class PaygenStageTest(generic_stages_unittest.AbstractStageTestCase,
                       cbuildbot_unittest.SimpleBuilderTestCase):
-  """Test the PaygenStageStage."""
+  """Test the PaygenStage Stage."""
 
   # We use a variant board to make sure the '_' is translated to '-'.
   BOT_ID = 'x86-alex_he-release'
@@ -362,23 +362,6 @@
       # Verify that we queue up work
       self.assertEqual(queue.put.call_args_list, [])
 
-  def testPerformStageBackgroundFail(self):
-    """Test that exception from background processes are properly handled."""
-    with patch(paygen_build_lib, 'CreatePayloads') as create_payloads:
-      create_payloads.side_effect = failures_lib.TestLabFailure
-
-      stage = self.ConstructStage(channels=['foo', 'bar'])
-
-      with patch(stage, '_HandleExceptionAsWarning') as warning_handler:
-        warning_handler.return_value = (results_lib.Results.FORGIVEN,
-                                        'description',
-                                        0)
-
-        stage.Run()
-
-        # This proves the exception was turned into a warning.
-        self.assertTrue(warning_handler.called)
-
   def testPerformStageTrybot(self):
     """Test the PerformStage alternate behavior for trybot runs."""
     with patch(release_stages.parallel, 'BackgroundTaskRunner') as background:
@@ -414,10 +397,23 @@
   def testRunPaygenInProcess(self):
     """Test that _RunPaygenInProcess works in the simple case."""
     with patch(paygen_build_lib, 'CreatePayloads') as create_payloads:
-      # Call the method under test.
+      # Have to patch and verify that the PaygenTestStage is created.
       stage = self.ConstructStage()
-      stage._RunPaygenInProcess('foo', 'foo-board', 'foo-version',
-                                False, False, False, skip_duts_check=False)
+      # CreatePayloads should return a tuple of a suite name and the finished
+      # URI.
+      create_payloads.side_effect = iter([('foo-suite-name',
+                                           'foo-archive-board',
+                                           'foo-archive-build',
+                                           'foo-finished_uri')])
+      with patch(paygen_build_lib, 'ScheduleAutotestTests') as sched_tests:
+        # Call the method under test.
+        stage._RunPaygenInProcess('foo', 'foo-board', 'foo-version',
+                                  False, False, False, skip_duts_check=False)
+        # Ensure that PaygenTestStage is created and schedules the test suite
+        # with the correct arguments.)
+        sched_tests.assert_called_once_with(
+            'foo-suite-name', 'foo-archive-board', 'foo-archive-build',
+            False, False)
 
       # Ensure arguments are properly converted and passed along.
       create_payloads.assert_called_with(gspaths.Build(version='foo-version',
@@ -434,6 +430,13 @@
   def testRunPaygenInProcessComplex(self):
     """Test that _RunPaygenInProcess with arguments that are more unusual."""
     with patch(paygen_build_lib, 'CreatePayloads') as create_payloads:
+      # Make CreatePayloads return the suite_name, archive_board, archive_build,
+      # and finished_uri.
+      create_payloads.side_effect = iter([('foo-suite-name',
+                                           'foo-archive-board',
+                                           'foo-archive-build',
+                                           'foo-finished_uri')])
+
       # Call the method under test.
       # Use release tools channel naming, and a board name including a variant.
       stage = self.ConstructStage()
@@ -453,3 +456,47 @@
           skip_delta_payloads=True,
           disable_tests=True,
           skip_duts_check=False)
+
+
+class PaygenTestStageTest(generic_stages_unittest.AbstractStageTestCase,
+                          cbuildbot_unittest.SimpleBuilderTestCase):
+  """Test the PaygenTestStage stage."""
+
+  # We use a variant board to make sure the '_' is translated to '-'.
+  BOT_ID = 'x86-alex_he-release'
+  RELEASE_TAG = '0.0.1'
+
+  def setUp(self):
+    self._Prepare()
+
+    # This method fetches a file from GS, mock it out.
+    # self.validateMock = self.PatchObject(
+    #     paygen_build_lib, 'ValidateBoardConfig')
+
+  # pylint: disable=arguments-differ
+  def ConstructStage(self, suite_name, build, finished_uri='foo-finished-uri',
+                     skip_duts_check=False):
+    return release_stages.PaygenTestStage(suite_name, self._current_board,
+                                          build, finished_uri,
+                                          skip_duts_check,
+                                          self._run)
+
+  def testPerformStageTestLabFail(self):
+    """Test that exception from RunHWTestSuite are properly handled."""
+    with patch(paygen_build_lib, 'ScheduleAutotestTests') as sched_tests:
+      sched_tests.side_effect = failures_lib.TestLabFailure
+
+      stage = self.ConstructStage('foo-test-suite',
+                                  gspaths.Build(version='foo-version',
+                                                board=self._current_board,
+                                                channel='foo-channel'))
+
+      with patch(stage, '_HandleExceptionAsWarning') as warning_handler:
+        warning_handler.return_value = (results_lib.Results.FORGIVEN,
+                                        'description',
+                                        0)
+
+        stage.Run()
+
+        # This proves the exception was turned into a warning.
+        self.assertTrue(warning_handler.called)
diff --git a/lib/paygen/paygen_build_lib.py b/lib/paygen/paygen_build_lib.py
index 2d22c9b..8b7f36b 100644
--- a/lib/paygen/paygen_build_lib.py
+++ b/lib/paygen/paygen_build_lib.py
@@ -16,7 +16,6 @@
 import json
 import operator
 import os
-import socket
 import sys
 import tempfile
 import urlparse
@@ -1081,33 +1080,6 @@
     logging.info('Control file emitted at %s', control_file)
     return control_file
 
-  def _ScheduleAutotestTests(self, suite_name):
-    """Run the appropriate command to schedule the Autotests we have prepped.
-
-    Args:
-      suite_name: The name of the test suite.
-    """
-    timeout_mins = config_lib.HWTestConfig.SHARED_HW_TEST_TIMEOUT / 60
-    cmd_result = commands.RunHWTestSuite(
-        board=self._archive_board,
-        build=self._archive_build,
-        suite=suite_name,
-        file_bugs=True,
-        pool='bvt',
-        priority=constants.HWTEST_BUILD_PRIORITY,
-        retry=True,
-        wait_for_results=True,
-        timeout_mins=timeout_mins,
-        suite_min_duts=2,
-        debug=bool(self._drm),
-        skip_duts_check=self._skip_duts_check)
-    if cmd_result.to_raise:
-      if isinstance(cmd_result.to_raise, failures_lib.TestWarning):
-        logging.warning('Warning running test suite; error output:\n%s',
-                        cmd_result.to_raise)
-      else:
-        raise cmd_result.to_raise
-
   def _AutotestPayloads(self, payload_tests):
     """Create necessary test artifacts and initiate Autotest runs.
 
@@ -1169,8 +1141,10 @@
           'Failed to infer archive build milestone number (%s)' %
           self._archive_build)
 
-    # Actually have the tests run.
-    self._ScheduleAutotestTests(suite_name)
+    # Send the information needed to actually schedule and run the tests.
+    finished_uri = self._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
+    return suite_name, finished_uri
+
 
   @staticmethod
   def _IsTestDeltaPayload(payload):
@@ -1299,6 +1273,7 @@
     """
     lock_uri = self._GetFlagURI(gspaths.ChromeosReleases.LOCK)
     finished_uri = self._GetFlagURI(gspaths.ChromeosReleases.FINISHED)
+    suite_name = None
 
     logging.info('Examining: %s', self._build)
 
@@ -1377,27 +1352,26 @@
             self._archive_board = archive_board
             self._archive_build = archive_build
             self._archive_build_uri = archive_build_uri
-
-            # We have a control file directory and all payloads have been
-            # generated. Lets create the list of tests to conduct.
-            payload_tests = self._CreatePayloadTests(payload_manager)
-            if payload_tests:
-              logging.info('Initiating %d payload tests', len(payload_tests))
-              self._drm(self._AutotestPayloads, payload_tests)
           except ArchiveError as e:
             logging.warning('Cannot map build to images archive, skipping '
                             'payload autotesting.')
             can_finish = False
 
+          if can_finish:
+            # We have a control file directory and all payloads have been
+            # generated. Lets create the list of tests to conduct.
+            payload_tests = self._CreatePayloadTests(payload_manager)
+            if payload_tests:
+              logging.info('Uploading %d payload tests', len(payload_tests))
+              suite_name, finished_uri = self._drm(self._AutotestPayloads,
+                                                   payload_tests)
+
         self._CleanupBuild()
-        if can_finish:
-          self._drm(gslib.CreateWithContents, finished_uri,
-                    socket.gethostname())
-        else:
+        if not can_finish:
           logging.warning('Not all payloads were generated, uploaded or '
                           'tested; not marking build as finished')
 
-        logging.info('Finished: %s', self._build)
+        logging.info('Finished generating payloads: %s', self._build)
 
     except gslock.LockNotAcquired as e:
       logging.info('Build already being processed: %s', e)
@@ -1411,6 +1385,8 @@
       logging.error('Failed: %s', self._build)
       raise
 
+    return suite_name, self._archive_board, self._archive_build, finished_uri
+
 
 def _FindControlFileDir(work_dir):
   """Decide the directory for emitting control files.
@@ -1475,14 +1451,47 @@
   """
   ValidateBoardConfig(build.board)
 
-  _PaygenBuild(build, work_dir, site_config,
-               dry_run=dry_run,
-               ignore_finished=ignore_finished,
-               skip_delta_payloads=skip_delta_payloads,
-               disable_tests=disable_tests,
-               output_dir=output_dir,
-               run_parallel=run_parallel,
-               skip_duts_check=skip_duts_check).CreatePayloads()
+  return _PaygenBuild(build, work_dir, site_config,
+                      dry_run=dry_run,
+                      ignore_finished=ignore_finished,
+                      skip_delta_payloads=skip_delta_payloads,
+                      disable_tests=disable_tests,
+                      output_dir=output_dir,
+                      run_parallel=run_parallel,
+                      skip_duts_check=skip_duts_check).CreatePayloads()
+
+def ScheduleAutotestTests(suite_name, board, build, skip_duts_check,
+                          debug):
+  """Run the appropriate command to schedule the Autotests we have prepped.
+
+  Args:
+  suite_name: The name of the test suite.
+  board: A string representing the name of the archive board.
+  build: A string representing the name of the archive build.
+  skip_duts_check: A boolean indicating to not check minimum available DUTs.
+  debug: A boolean indicating whether or not we are in debug mode.
+  """
+  timeout_mins = config_lib.HWTestConfig.SHARED_HW_TEST_TIMEOUT / 60
+  cmd_result = commands.RunHWTestSuite(
+      board=board,
+      build=build,
+      suite=suite_name,
+      file_bugs=True,
+      pool='bvt',
+      priority=constants.HWTEST_BUILD_PRIORITY,
+      retry=True,
+      wait_for_results=True,
+      timeout_mins=timeout_mins,
+      suite_min_duts=2,
+      debug=debug,
+      skip_duts_check=skip_duts_check)
+
+  if cmd_result.to_raise:
+    if isinstance(cmd_result.to_raise, failures_lib.TestWarning):
+      logging.warning('Warning running test suite; error output:\n%s',
+                      cmd_result.to_raise)
+    else:
+      raise cmd_result.to_raise
 
 
 def _GetAndValidateJson(uri):
diff --git a/lib/paygen/paygen_build_lib_unittest.py b/lib/paygen/paygen_build_lib_unittest.py
index eac297b..640be30 100644
--- a/lib/paygen/paygen_build_lib_unittest.py
+++ b/lib/paygen/paygen_build_lib_unittest.py
@@ -1175,10 +1175,9 @@
     paygen._MapToArchive('foo-board', '1.2.3').AndReturn(
         ('archive_board', 'archive_build', 'archive_build_uri'))
     paygen._CreatePayloadTests(payload_manager).AndReturn(['Test Payloads'])
-    paygen._AutotestPayloads(['Test Payloads'])
-
+    paygen._AutotestPayloads(['Test Payloads']).AndReturn(('foo-suite-name',
+                                                           finished_uri))
     paygen._CleanupBuild()
-    gslib.CreateWithContents(finished_uri, mox.IgnoreArg())
     lock.__exit__(
         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
 
@@ -1215,8 +1214,8 @@
     paygen._MapToArchive('foo-board', '1.2.3').AndReturn(
         ('archive_board', 'archive_build', 'archive_build_uri'))
     paygen._CreatePayloadTests(payload_manager).AndReturn(['Test Payloads'])
-    paygen._AutotestPayloads(['Test Payloads'])
-    gslib.CreateWithContents(finished_uri, mox.IgnoreArg())
+    paygen._AutotestPayloads(['Test Payloads']).AndReturn(('foo-suite-name',
+                                                           finished_uri))
     paygen._CleanupBuild()
     lock.__exit__(
         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
@@ -1248,7 +1247,6 @@
     paygen_payload_lib.FindExistingPayloads(payload).AndReturn([])
     paygen._GeneratePayloads(payload_list, lock)
     paygen._CleanupBuild()
-    gslib.CreateWithContents(finished_uri, mox.IgnoreArg())
     lock.__exit__(
         mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(None)
 
@@ -1257,6 +1255,30 @@
 
     paygen.CreatePayloads()
 
+  def setupAutotestPayloadsTests(self):
+    paygen = self._GetPaygenBuildInstance()
+
+    self.mox.StubOutWithMock(paygen, '_EmitControlFile')
+
+    return paygen
+
+  def testAutotestPayloads(self):
+    """Test that Autotest tarballs are uploaded and metadata returned."""
+    paygen = self.setupAutotestPayloadsTests()
+
+    payload_tests = ['Payload Tests']
+    paygen._archive_build_uri = 'archive_build_uri'
+    paygen._archive_build = 'foo-bd-release/R9001-1234.56.7'
+
+    test_channel = paygen._build.channel.rpartition('-')[0]
+    suite_name = (paygen.PAYGEN_AU_SUITE_TEMPLATE % test_channel)
+    for test in payload_tests:
+      paygen._EmitControlFile(test, suite_name, mox.IgnoreArg())
+
+    self.mox.ReplayAll()
+
+    paygen._AutotestPayloads(payload_tests)
+
   def setupCreatePayloadTests(self):
     paygen = self._GetPaygenBuildInstance()
 
@@ -1459,10 +1481,6 @@
 
   def testScheduleAutotestTests(self):
     """Test scheduling autotest tests with build autotest proxy."""
-    paygen = paygen_build_lib._PaygenBuild(
-        self.foo_build, self.tempdir,
-        config_lib_unittest.MockSiteConfig())
-
     self.mox.StubOutWithMock(commands, 'RunHWTestSuite')
     self.mox.StubOutWithMock(cros_build_lib, 'RunCommand')
 
@@ -1478,17 +1496,18 @@
     self.mox.ReplayAll()
 
     # Setup preliminary values needed for scheduling autotests.
-    paygen._archive_board = 'foo-board'
-    paygen._archive_build = 'foo-board-release/R99-1.2.3'
+    suite_name = 'paygen_au_foo'
+    archive_board = 'foo-board'
+    archive_build = 'foo-board-release/R99-1.2.3'
+    skip_duts_check = False
+    debug = False
 
-    paygen._ScheduleAutotestTests('paygen_au_foo')
+    paygen_build_lib.ScheduleAutotestTests(suite_name, archive_board,
+                                           archive_build, skip_duts_check,
+                                           debug)
 
   def testScheduleAutotestTestsWarn(self):
     """Test scheduling autotest tests with build autotest proxy."""
-    paygen = paygen_build_lib._PaygenBuild(
-        self.foo_build, self.tempdir,
-        config_lib_unittest.MockSiteConfig())
-
     self.mox.StubOutWithMock(commands, 'RunHWTestSuite')
     self.mox.StubOutWithMock(cros_build_lib, 'RunCommand')
 
@@ -1506,10 +1525,15 @@
     self.mox.ReplayAll()
 
     # Setup preliminary values needed for scheduling autotests.
-    paygen._archive_board = 'foo-board'
-    paygen._archive_build = 'foo-board-release/R99-1.2.3'
+    suite_name = 'paygen_au_foo'
+    archive_board = 'foo-board'
+    archive_build = 'foo-board-release/R99-1.2.3'
+    skip_duts_check = False
+    debug = False
 
-    paygen._ScheduleAutotestTests('paygen_au_foo')
+    paygen_build_lib.ScheduleAutotestTests(suite_name, archive_board,
+                                           archive_build, skip_duts_check,
+                                           debug)
 
   def testMapToArchive(self):
     """Test that mapping to images archive names/locations works."""
@@ -1525,8 +1549,8 @@
     site_config.Add('build_to_introduce_boards',
                     boards=['foo_board', 'bar_board', 'bar-board'])
 
-    paygen = paygen_build_lib._PaygenBuild(
-        self.foo_build, self.tempdir, site_config)
+    paygen = paygen_build_lib._PaygenBuild(self.foo_build, self.tempdir,
+                                           site_config)
 
     # Case 1: mapping successful.
     self.assertEqual(