[infra] Improve test speed (#5118)

Make unittests take 20 seconds to run instead of 35.
Make integration tests take 50 seconds to run instead of 6 minutes.
Make CI take 6 minutes instead of 12 minutes.

1. Allow running tests in parallel. Locally this takes the time for running all tests (including integration tests) from 6 minutes to ~50 seconds. We don't do parallel by default since it doesn't really save any time unless running integration tests on my machine (probably due to overhead of starting ~70 processes). This also speeds up CI from about 12 minutes to 6 minutes  (since github actions has 2 cores per machine).
2. Fix how we run tests. I'm not exactly sure why, but the method we used for discovering tests, recursing through every directory and passing to unittest caused the build/infra tests to execute twice. Fixing this makes running unittests take ~20 seconds instead of ~35.

This change also uses pytest for running tests since it's easy to use it to run tests in parallel.
This change was made possible by #5113

4 files changed, 51 insertions, 24 deletions

diff --git a/infra/ci/build_test.py b/infra/ci/build_test.py
index 59fa3cf5c..9da67e253 100644
--- a/infra/ci/build_test.py
+++ b/infra/ci/build_test.py
@@ -16,9 +16,14 @@
 """Tests for build.py"""
 
 import os
+import sys
 import unittest
 from unittest import mock
 
+# pylint: disable=wrong-import-position
+INFRA_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(INFRA_DIR)
+
 from ci import build
 
 
diff --git a/infra/ci/requirements.txt b/infra/ci/requirements.txt
index 2a21c2af5..48d2ae758 100644
--- a/infra/ci/requirements.txt
+++ b/infra/ci/requirements.txt
@@ -2,5 +2,7 @@
 parameterized==0.7.4
 pyfakefs==4.1.0
 pylint==2.5.3
+pytest==6.2.1
+pytest-xdist==2.2.0
 PyYAML==5.3.1
 yapf==0.30.0
diff --git a/infra/presubmit.py b/infra/presubmit.py
index a67ba0b7f..bd5e9c327 100755
--- a/infra/presubmit.py
+++ b/infra/presubmit.py
@@ -18,17 +18,12 @@
 
 import argparse
 import os
-import re
 import subprocess
 import sys
 import unittest
 import yaml
 
 _SRC_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-TEST_BLOCKLIST = [
-    # Test errors with error: "ModuleNotFoundError: No module named 'apt'".
-    re.compile(r'.*\/infra\/base-images\/base-sanitizer-libs-builder'),
-]
 
 
 def _is_project_file(actual_path, expected_filename):
@@ -345,32 +340,49 @@ def get_changed_files():
   return [os.path.abspath(f) for f in changed_files]
 
 
-def is_test_dir_blocklisted(directory):
-  """Returns True if |directory| is blocklisted."""
-  for blocklisted_regex in TEST_BLOCKLIST:
-    if blocklisted_regex.search(directory):
-      return True
-  return False
+def run_build_tests():
+  """Runs build tests because they can't be run in parallel."""
+  suite_list = [
+      unittest.TestLoader().discover(os.path.join(_SRC_ROOT, 'infra', 'build'),
+                                     pattern='*_test.py'),
+  ]
+  suite = unittest.TestSuite(suite_list)
+  print('Running build tests.')
+  result = unittest.TextTestRunner().run(suite)
+  return not result.failures and not result.errors
 
 
-def run_tests(_=None):
-  """Run all unit tests."""
+def run_nonbuild_tests(parallel):
+  """Run all tests but build tests. Do it in parallel if |parallel|. The reason
+  why we exclude build tests is because they use an emulator that prevents them
+  from being used in parallel."""
+  # We look for all project directories because otherwise pytest won't run tests
+  # that are not in valid modules (e.g. "base-images").
   relevant_dirs = set()
   all_files = get_all_files()
   for file_path in all_files:
     directory = os.path.dirname(file_path)
-    if is_test_dir_blocklisted(directory):
-      continue
     relevant_dirs.add(directory)
 
-  suite_list = []
-  for relevant_dir in relevant_dirs:
-    suite_list.append(unittest.TestLoader().discover(relevant_dir,
-                                                     pattern='*_test.py'))
-  suite = unittest.TestSuite(suite_list)
-  result = unittest.TextTestRunner().run(suite)
+  # Use ignore-glob because ignore doesn't seem to work properly with the way we
+  # pass directories to pytest.
+  command = [
+      'pytest',
+      # Test errors with error: "ModuleNotFoundError: No module named 'apt'.
+      '--ignore-glob=infra/base-images/base-sanitizer-libs-builder/*',
+      '--ignore-glob=infra/build/*',
+  ]
+  if parallel:
+    command.extend(['-n', 'auto'])
+  command += list(relevant_dirs)
+  print('Running non-build tests.')
+  return subprocess.run(command, check=False).returncode == 0
 
-  return not result.failures and not result.errors
+
+def run_tests(_=None, parallel=False):
+  """Runs all unit tests."""
+  success = run_nonbuild_tests(parallel)
+  return success and run_build_tests()
 
 
 def get_all_files():
@@ -387,10 +399,16 @@ def main():
   parser.add_argument('command',
                       choices=['format', 'lint', 'license', 'infra-tests'],
                       nargs='?')
-  parser.add_argument('--all-files',
+  parser.add_argument('-a',
+                      '--all-files',
                       action='store_true',
                       help='Run presubmit check(s) on all files',
                       default=False)
+  parser.add_argument('-p',
+                      '--parallel',
+                      action='store_true',
+                      help='Run tests in parallel.',
+                      default=False)
   args = parser.parse_args()
 
   if args.all_files:
@@ -414,7 +432,7 @@ def main():
     return bool_to_returncode(success)
 
   if args.command == 'infra-tests':
-    success = run_tests(relevant_files)
+    success = run_tests(relevant_files, parallel=args.parallel)
     return bool_to_returncode(success)
 
   # Do all the checks (but no tests).
diff --git a/infra/pytest.ini b/infra/pytest.ini
new file mode 100644
index 000000000..d9bb3737e
--- /dev/null
+++ b/infra/pytest.ini
@@ -0,0 +1,2 @@
+[pytest]
+python_files = *_test.py
\ No newline at end of file