[autotest] make build_externals support testing for chromite

In order for testing chromite during push, build_externals should
support updating chromite repo to master branch

BUG=None
TEST=
$./utils/build_externals.py
 # this command updates chromite to prod branch successfully
$./utils/build_externals.py --use_chromite_master
 # this command updates chromite to master branch successfully
$ ./site_utils/deploy_server_local.py
 # this command updates chromite to prod branch successfully
$./site_utils/deploy_server_local.py --update_push_servers
 # this command updates chromite to master branch successfully

Change-Id: I29196084943d11984f5e9cc61693d57e9e148caf
Reviewed-on: https://chromium-review.googlesource.com/401180
Reviewed-by: Shuqian Zhao <shuqianz@chromium.org>
Commit-Queue: Shuqian Zhao <shuqianz@chromium.org>
Tested-by: Shuqian Zhao <shuqianz@chromium.org>
diff --git a/site_utils/deploy_server_local.py b/site_utils/deploy_server_local.py
index ce28126..cc728d5 100755
--- a/site_utils/deploy_server_local.py
+++ b/site_utils/deploy_server_local.py
@@ -41,6 +41,7 @@
 # frontend repo, no need to update afe.
 COMMANDS_TO_REPOS_DICT = {'afe': 'frontend/',
                           'tko': 'tko/'}
+BUILD_EXTERNALS_COMMAND = 'build_externals'
 # Services present on all hosts.
 # TODO(ayatane): Temporarily stop starting sysmon
 # UNIVERSAL_SERVICES = ['sysmon']
@@ -192,7 +193,7 @@
     return services
 
 
-def update_command(cmd_tag, dryrun=False):
+def update_command(cmd_tag, dryrun=False, use_chromite_master=False):
     """Restart a command.
 
     The command name is looked up in global_config.ini to find the full command
@@ -200,6 +201,8 @@
 
     @param cmd_tag: Which command to restart.
     @param dryrun: If true print the command that would have been run.
+    @param use_chromite_master: True if updating chromite to master, rather
+                                than prod.
 
     @raises UnknownCommandException If cmd_tag can't be looked up.
     @raises subprocess.CalledProcessError on a command failure.
@@ -214,6 +217,10 @@
 
     expanded_command = cmds[cmd_tag].replace('AUTOTEST_REPO',
                                               common.autotest_dir)
+    # When updating push servers, pass an arg to build_externals to update
+    # chromite to master branch for testing
+    if use_chromite_master and cmd_tag == BUILD_EXTERNALS_COMMAND:
+        expanded_command += ' --use_chromite_master'
 
     print('Running: %s: %s' % (cmd_tag, expanded_command))
     if dryrun:
@@ -308,7 +315,7 @@
 
 
 def run_deploy_actions(cmds_skip=set(), dryrun=False,
-                       skip_service_status=False):
+                       skip_service_status=False, use_chromite_master=False):
     """Run arbitrary update commands specified in global.ini.
 
     @param cmds_skip: cmds no need to run since the corresponding repo/file
@@ -316,6 +323,8 @@
     @param dryrun: Don't really restart the service, just print out the command.
     @param skip_service_status: Set to True to skip service status check.
                                 Default is False.
+    @param use_chromite_master: True if updating chromite to master, rather
+                                than prod.
 
     @raises subprocess.CalledProcessError on a command failure.
     @raises UnstableServices if any services are unstable after restart.
@@ -330,7 +339,8 @@
                             status='primary')):
                 print('Command %s is only applicable to primary servers.' % cmd)
                 continue
-            update_command(cmd, dryrun=dryrun)
+            update_command(cmd, dryrun=dryrun,
+                           use_chromite_master=use_chromite_master)
 
     services = discover_restart_services()
     if services:
@@ -476,9 +486,9 @@
             return 1
 
     versions_before = repo_versions()
-    versions_after = {}
+    versions_after = set()
     cmd_versions_before = repo_versions_to_decide_whether_run_cmd_update()
-    cmd_versions_after = {}
+    cmd_versions_after = set()
 
     if behaviors.update:
         print('Updating Repo.')
@@ -494,7 +504,8 @@
                      {t[0] for t in cmd_versions_before & cmd_versions_after})
         try:
             run_deploy_actions(
-                    cmds_skip, behaviors.dryrun, behaviors.skip_service_status)
+                    cmds_skip, behaviors.dryrun, behaviors.skip_service_status,
+                    use_chromite_master=behaviors.update_push_servers)
         except UnstableServices as e:
             print('The following services were not stable after '
                   'the update:')
diff --git a/utils/build_externals.py b/utils/build_externals.py
index ccd102a..7841e83 100755
--- a/utils/build_externals.py
+++ b/utils/build_externals.py
@@ -12,7 +12,12 @@
     utils/build_externals.py
 """
 
-import compileall, logging, os, sys
+import argparse
+import compileall
+import logging
+import os
+import sys
+
 import common
 from autotest_lib.client.common_lib import logging_config, logging_manager
 from autotest_lib.client.common_lib import utils
@@ -47,6 +52,7 @@
     Find all ExternalPackage classes defined in this file and ask them to
     fetch, build and install themselves.
     """
+    options = parse_arguments(sys.argv[1:])
     logging_manager.configure_logging(BuildExternalsLoggingConfig(),
                                       verbose=True)
     os.umask(022)
@@ -66,9 +72,10 @@
         os.environ[env_python_path_varname] = ':'.join([
             install_dir, env_python_path])
 
-    fetched_packages, fetch_errors = fetch_necessary_packages(package_dir,
-                                                              install_dir)
-    install_errors = build_and_install_packages(fetched_packages, install_dir)
+    fetched_packages, fetch_errors = fetch_necessary_packages(
+        package_dir, install_dir, set(options.names_to_check))
+    install_errors = build_and_install_packages(
+        fetched_packages, install_dir, options.use_chromite_master)
 
     # Byte compile the code after it has been installed in its final
     # location as .pyc files contain the path passed to compile_dir().
@@ -88,19 +95,39 @@
     return len(errors)
 
 
-def fetch_necessary_packages(dest_dir, install_dir):
+def parse_arguments(args):
+    """Parse command line arguments.
+
+    @param args: The command line arguments to parse. (ususally sys.argsv[1:])
+
+    @returns An argparse.Namespace populated with argument values.
+    """
+    parser = argparse.ArgumentParser(
+            description='Command to build third party dependencies required '
+                        'for autotest.')
+    parser.add_argument('--use_chromite_master', action='store_true',
+                        help='Update chromite to master branch, rather than '
+                             'prod.')
+    parser.add_argument('--names_to_check', nargs='*', type=str, default=set(),
+                        help='Package names to check whether they are needed '
+                             'in current system.')
+    return parser.parse_args(args)
+
+
+def fetch_necessary_packages(dest_dir, install_dir, names_to_check=set()):
     """
     Fetches all ExternalPackages into dest_dir.
 
     @param dest_dir: Directory the packages should be fetched into.
     @param install_dir: Directory where packages will later installed.
+    @param names_to_check: A set of package names to check whether they are
+                           needed on current system. Default is empty.
 
     @returns A tuple containing two lists:
              * A list of ExternalPackage instances that were fetched and
                need to be installed.
              * A list of error messages for any failed fetches.
     """
-    names_to_check = sys.argv[1:]
     errors = []
     fetched_packages = []
     for package_class in external_packages.ExternalPackage.subclasses:
@@ -124,18 +151,23 @@
     return fetched_packages, errors
 
 
-def build_and_install_packages(packages, install_dir):
+def build_and_install_packages(packages, install_dir,
+                               use_chromite_master=False):
     """
     Builds and installs all packages into install_dir.
 
     @param packages - A list of already fetched ExternalPackage instances.
     @param install_dir - Directory the packages should be installed into.
+    @param use_chromite_master: True if updating chromite to master branch.
 
     @returns A list of error messages for any installs that failed.
     """
     errors = []
     for package in packages:
-        result = package.build_and_install(install_dir)
+        if use_chromite_master and package.name.lower() == 'chromiterepo':
+            result = package.build_and_install(install_dir, master_branch=True)
+        else:
+            result = package.build_and_install(install_dir)
         if isinstance(result, bool):
             success = result
             message = None
diff --git a/utils/external_packages.py b/utils/external_packages.py
index f25cbe4..fbbce30 100644
--- a/utils/external_packages.py
+++ b/utils/external_packages.py
@@ -1173,6 +1173,7 @@
     # All the chromiumos projects used on the lab servers should have a 'prod'
     # branch used to track the software version deployed in prod.
     PROD_BRANCH = 'prod'
+    MASTER_BRANCH = 'master'
 
     def is_needed(self, unused_install_dir):
         """Tell build_externals that we need to fetch."""
@@ -1246,7 +1247,7 @@
 
     _GIT_URL = ('https://chromium.googlesource.com/chromiumos/chromite')
 
-    def build_and_install(self, install_dir):
+    def build_and_install(self, install_dir, master_branch=False):
         """
         Clone if the repo isn't initialized, pull clean bits if it is.
 
@@ -1255,13 +1256,17 @@
         directory because it doesn't need installation.
 
         @param install_dir: destination directory for chromite installation.
+        @param master_branch: if True, install master branch. Otherwise,
+                              install prod branch.
         """
+        init_branch = (self.MASTER_BRANCH if master_branch
+                       else self.PROD_BRANCH)
         local_chromite_dir = os.path.join(install_dir, 'chromite')
         git_repo = revision_control.GitRepo(
                 local_chromite_dir,
                 self._GIT_URL,
                 abs_work_tree=local_chromite_dir)
-        git_repo.reinit_repo_at(self.PROD_BRANCH)
+        git_repo.reinit_repo_at(init_branch)
 
 
         if git_repo.get_latest_commit_hash():