[trial_build] Fix bugs (#7764)

* Fix bugs

1. Don't try to build the script name as a project.
2. Add a flag to force builds of projects that previously failed.
3. Make sure we build projects from our PR branch.
4. Wait on all builds not just builds of the last type.
5. Don't use test bucket for corpus or coverage (will corpus work or fail because of creds?)

Add tests for these features.

9 files changed, 533 insertions, 35 deletions

diff --git a/infra/build/functions/build_and_run_coverage.py b/infra/build/functions/build_and_run_coverage.py
index 23a15a87..48c89c82 100755
--- a/infra/build/functions/build_and_run_coverage.py
+++ b/infra/build/functions/build_and_run_coverage.py
@@ -314,8 +314,7 @@ def get_fuzz_introspector_steps(  # pylint: disable=too-many-locals, too-many-ar
                   f'/reports/{coverage_report_latest}/linux')
 
   download_coverage_steps = build_lib.download_coverage_data_steps(
-      project.name, coverage_report_latest, bucket_name, build.out,
-      config.testing)
+      project.name, coverage_report_latest, bucket_name, build.out)
   if not download_coverage_steps:
     logging.warning(
         'Skipping introspector build for %s. No coverage data found.',
diff --git a/infra/build/functions/build_lib.py b/infra/build/functions/build_lib.py
index f1e46dc5..52567525 100644
--- a/infra/build/functions/build_lib.py
+++ b/infra/build/functions/build_lib.py
@@ -105,11 +105,13 @@ def get_upload_bucket(engine, architecture, testing):
   return bucket
 
 
-def _get_targets_list(project_name, testing):
+def _get_targets_list(project_name):
   """Returns target list."""
   # libFuzzer ASan 'x86_84' is the default configuration, get list of targets
   # from it.
-  bucket = get_upload_bucket('libfuzzer', 'x86_64', testing)
+  # We never want the target list from the testing bucket, the testing bucket is
+  # only for uploading.
+  bucket = get_upload_bucket('libfuzzer', 'x86_64', testing=None)
   url = get_targets_list_url(bucket, project_name, 'address')
 
   url = urlparse.urljoin(GCS_URL_BASENAME, url)
@@ -159,10 +161,10 @@ def get_signed_url(path, method='PUT', content_type=''):
   return f'https://storage.googleapis.com{path}?{urlparse.urlencode(values)}'
 
 
-def download_corpora_steps(project_name, testing):
+def download_corpora_steps(project_name):
   """Returns GCB steps for downloading corpora backups for the given project.
   """
-  fuzz_targets = _get_targets_list(project_name, testing)
+  fuzz_targets = _get_targets_list(project_name)
   if not fuzz_targets:
     sys.stderr.write('No fuzz targets found for project "%s".\n' % project_name)
     return None
@@ -197,11 +199,10 @@ def download_corpora_steps(project_name, testing):
   return steps
 
 
-def download_coverage_data_steps(project_name, latest, bucket_name, out_dir,
-                                 testing):
+def download_coverage_data_steps(project_name, latest, bucket_name, out_dir):
   """Returns GCB steps to download coverage data for the given project"""
   steps = []
-  fuzz_targets = _get_targets_list(project_name, testing)
+  fuzz_targets = _get_targets_list(project_name)
   if not fuzz_targets:
     sys.stderr.write('No fuzz targets found for project "%s".\n' % project_name)
     return None
diff --git a/infra/build/functions/build_project.py b/infra/build/functions/build_project.py
index 52e76cd4..01523bdd 100755
--- a/infra/build/functions/build_project.py
+++ b/infra/build/functions/build_project.py
@@ -274,8 +274,11 @@ def get_build_steps(  # pylint: disable=too-many-locals, too-many-statements, to
   # Sort engines to make AFL first to test if libFuzzer has an advantage in
   # finding bugs first since it is generally built first.
   for fuzzing_engine in sorted(project.fuzzing_engines):
-    for sanitizer in project.sanitizers:
-      for architecture in project.architectures:
+    # Sort sanitizers and architectures so order is determinisitic (good for
+    # tests).
+    for sanitizer in sorted(project.sanitizers):
+      # Build x86_64 before i386.
+      for architecture in reversed(sorted(project.architectures)):
         build = Build(fuzzing_engine, sanitizer, architecture)
         if not is_supported_configuration(build):
           continue
diff --git a/infra/build/functions/ci_trial_build.py b/infra/build/functions/ci_trial_build.py
index e829742e..0ad8ede3 100644
--- a/infra/build/functions/ci_trial_build.py
+++ b/infra/build/functions/ci_trial_build.py
@@ -26,7 +26,7 @@ import github
 
 import trial_build
 
-TRIGGER_COMMAND = '/gcbrun '
+TRIGGER_COMMAND = '/gcbrun'
 
 
 def get_comments(pull_request_number):
@@ -48,23 +48,29 @@ def get_latest_gcbrun_command(comments):
   for comment in reversed(comments):
     # This seems to get comments on code too.
     body = comment.body
-    if not body.startswith(TRIGGER_COMMAND):
+    command_str = f'{TRIGGER_COMMAND} trial_build.py '
+    if not body.startswith(command_str):
       continue
-    if len(body) <= len(TRIGGER_COMMAND):
+    if len(body) == len(command_str):
       return None
-    # Add an extra for white space.
-    return body[len(TRIGGER_COMMAND):].strip().split(' ')
+    return body[len(command_str):].strip().split(' ')
   return None
 
 
-def exec_command_from_github(pull_request_number):
+def exec_command_from_github(pull_request_number, branch):
   """Executes the gcbrun command for trial_build.py in the most recent command
   on |pull_request_number|."""
   comments = get_comments(pull_request_number)
   command = get_latest_gcbrun_command(comments)
+  logging.info('Command: %s.', command)
   if command is None:
     logging.info('Trial build not requested.')
     return None
+
+  # Set the branch so that the trial_build builds the projects from the PR
+  # branch.
+  command.append(f' --branch {branch}')
+
   logging.info('Command: %s.', command)
   return trial_build.trial_build_main(command, local_base_build=False)
 
@@ -73,7 +79,8 @@ def main():
   """Entrypoint for GitHub CI into trial_build.py"""
   logging.basicConfig(level=logging.INFO)
   pull_request_number = int(os.environ['PULL_REQUEST_NUMBER'])
-  return 0 if exec_command_from_github(pull_request_number) else 1
+  branch = os.environ['BRANCH']
+  return 0 if exec_command_from_github(pull_request_number, branch) else 1
 
 
 if __name__ == '__main__':
diff --git a/infra/build/functions/ci_trial_build_test.py b/infra/build/functions/ci_trial_build_test.py
new file mode 100644
index 00000000..701977a1
--- /dev/null
+++ b/infra/build/functions/ci_trial_build_test.py
@@ -0,0 +1,54 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+################################################################################
+"""Tests for ci_trial_build.py."""
+import unittest
+from unittest import mock
+
+import ci_trial_build
+
+
+class GetLatestGCBrunCommandTest(unittest.TestCase):
+  """Tests for get_latest_gcbrun_command."""
+
+  def test_command_parsing(self):
+    """Tests that commands from GitHub comments are parsed properly."""
+    mock_comment = mock.MagicMock()
+    mock_comment.body = ('/gcbrun trial_build.py aiohttp --sanitizer '
+                         'coverage address --fuzzing-engine libfuzzer')
+    comments = [mock_comment]
+    expected_command = [
+        'aiohttp', '--sanitizer', 'coverage', 'address', '--fuzzing-engine',
+        'libfuzzer'
+    ]
+    actual_command = ci_trial_build.get_latest_gcbrun_command(comments)
+    self.assertEqual(expected_command, actual_command)
+
+  def test_last_comment(self):
+    """Tests that the last comment from the GitHub PR is considered the
+    command."""
+    mock_comment_1 = mock.MagicMock()
+    mock_comment_1.body = ('/gcbrun trial_build.py aiohttp --sanitizer '
+                           'coverage address --fuzzing-engine libfuzzer')
+    mock_comment_2 = mock.MagicMock()
+    mock_comment_2.body = ('/gcbrun trial_build.py skcms --sanitizer '
+                           'coverage address --fuzzing-engine libfuzzer')
+    comments = [mock_comment_1, mock_comment_2]
+    expected_command = [
+        'skcms', '--sanitizer', 'coverage', 'address', '--fuzzing-engine',
+        'libfuzzer'
+    ]
+    actual_command = ci_trial_build.get_latest_gcbrun_command(comments)
+    self.assertEqual(expected_command, actual_command)
diff --git a/infra/build/functions/test_data/expected_trial_build_steps.json b/infra/build/functions/test_data/expected_trial_build_steps.json
new file mode 100644
index 00000000..75259828
--- /dev/null
+++ b/infra/build/functions/test_data/expected_trial_build_steps.json
@@ -0,0 +1,332 @@
+[
+  {
+    "args": [
+      "clone",
+      "https://github.com/google/oss-fuzz.git",
+      "--depth",
+      "1",
+      "--branch",
+      "mybranch"
+    ],
+    "name": "gcr.io/cloud-builders/git"
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "pull",
+      "gcr.io/oss-fuzz-base/base-builder-testing"
+    ],
+    "waitFor": "-"
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "tag",
+      "gcr.io/oss-fuzz-base/base-builder-testing",
+      "gcr.io/oss-fuzz-base/base-builder"
+    ]
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "pull",
+      "gcr.io/oss-fuzz-base/base-builder-swift-testing"
+    ],
+    "waitFor": "-"
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "tag",
+      "gcr.io/oss-fuzz-base/base-builder-swift-testing",
+      "gcr.io/oss-fuzz-base/base-builder-swift"
+    ]
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "pull",
+      "gcr.io/oss-fuzz-base/base-builder-jvm-testing"
+    ],
+    "waitFor": "-"
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "tag",
+      "gcr.io/oss-fuzz-base/base-builder-jvm-testing",
+      "gcr.io/oss-fuzz-base/base-builder-jvm"
+    ]
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "pull",
+      "gcr.io/oss-fuzz-base/base-builder-go-testing"
+    ],
+    "waitFor": "-"
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "tag",
+      "gcr.io/oss-fuzz-base/base-builder-go-testing",
+      "gcr.io/oss-fuzz-base/base-builder-go"
+    ]
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "pull",
+      "gcr.io/oss-fuzz-base/base-builder-python-testing"
+    ],
+    "waitFor": "-"
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "tag",
+      "gcr.io/oss-fuzz-base/base-builder-python-testing",
+      "gcr.io/oss-fuzz-base/base-builder-python"
+    ]
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "pull",
+      "gcr.io/oss-fuzz-base/base-builder-rust-testing"
+    ],
+    "waitFor": "-"
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "tag",
+      "gcr.io/oss-fuzz-base/base-builder-rust-testing",
+      "gcr.io/oss-fuzz-base/base-builder-rust"
+    ]
+  },
+  {
+    "name": "gcr.io/cloud-builders/docker",
+    "args": [
+      "build",
+      "--tag",
+      "gcr.io/oss-fuzz/skcms",
+      "."
+    ],
+    "dir": "oss-fuzz/projects/skcms"
+  },
+  {
+    "name": "gcr.io/oss-fuzz/skcms",
+    "args": [
+      "bash",
+      "-c",
+      "srcmap > /workspace/srcmap.json && cat /workspace/srcmap.json"
+    ],
+    "env": [
+      "OSSFUZZ_REVISION=$REVISION_ID",
+      "FUZZING_LANGUAGE=c++"
+    ],
+    "id": "srcmap"
+  },
+  {
+    "name": "gcr.io/oss-fuzz/skcms",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=afl",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/afl-address-x86_64",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "rm -r /out && cd /src && cd skcms && mkdir -p /workspace/out/afl-address-x86_64 && compile || (echo \"********************************************************************************\nFailed to build.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer address --engine afl --architecture x86_64 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "compile-afl-address-x86_64"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=afl",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/afl-address-x86_64",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "test_all.py || (echo \"********************************************************************************\nBuild checks failed.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer address --engine afl --architecture x86_64 skcms\npython infra/helper.py check_build --sanitizer address --engine afl --architecture x86_64 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "build-check-afl-address-x86_64"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=afl",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/afl-address-x86_64",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "targets_list > /workspace/targets.list.address"
+    ]
+  },
+  {
+    "name": "gcr.io/oss-fuzz/skcms",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-address-x86_64",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "rm -r /out && cd /src && cd skcms && mkdir -p /workspace/out/libfuzzer-address-x86_64 && compile || (echo \"********************************************************************************\nFailed to build.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer address --engine libfuzzer --architecture x86_64 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "compile-libfuzzer-address-x86_64"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-address-x86_64",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "test_all.py || (echo \"********************************************************************************\nBuild checks failed.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer address --engine libfuzzer --architecture x86_64 skcms\npython infra/helper.py check_build --sanitizer address --engine libfuzzer --architecture x86_64 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "build-check-libfuzzer-address-x86_64"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-address-x86_64",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "targets_list > /workspace/targets.list.address"
+    ]
+  },
+  {
+    "name": "gcr.io/oss-fuzz/skcms",
+    "env": [
+      "ARCHITECTURE=i386",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-address-i386",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "rm -r /out && cd /src && cd skcms && mkdir -p /workspace/out/libfuzzer-address-i386 && compile || (echo \"********************************************************************************\nFailed to build.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer address --engine libfuzzer --architecture i386 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "compile-libfuzzer-address-i386"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=i386",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-address-i386",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "test_all.py || (echo \"********************************************************************************\nBuild checks failed.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer address --engine libfuzzer --architecture i386 skcms\npython infra/helper.py check_build --sanitizer address --engine libfuzzer --architecture i386 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "build-check-libfuzzer-address-i386"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=i386",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-address-i386",
+      "SANITIZER=address"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "targets_list > /workspace/targets.list.address"
+    ]
+  },
+  {
+    "name": "gcr.io/oss-fuzz/skcms",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-undefined-x86_64",
+      "SANITIZER=undefined"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "rm -r /out && cd /src && cd skcms && mkdir -p /workspace/out/libfuzzer-undefined-x86_64 && compile || (echo \"********************************************************************************\nFailed to build.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer undefined --engine libfuzzer --architecture x86_64 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "compile-libfuzzer-undefined-x86_64"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-undefined-x86_64",
+      "SANITIZER=undefined"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "test_all.py || (echo \"********************************************************************************\nBuild checks failed.\nTo reproduce, run:\npython infra/helper.py build_image skcms\npython infra/helper.py build_fuzzers --sanitizer undefined --engine libfuzzer --architecture x86_64 skcms\npython infra/helper.py check_build --sanitizer undefined --engine libfuzzer --architecture x86_64 skcms\n********************************************************************************\" && false)"
+    ],
+    "id": "build-check-libfuzzer-undefined-x86_64"
+  },
+  {
+    "name": "gcr.io/oss-fuzz-base/base-runner-testing",
+    "env": [
+      "ARCHITECTURE=x86_64",
+      "FUZZING_ENGINE=libfuzzer",
+      "FUZZING_LANGUAGE=c++",
+      "HOME=/root",
+      "OUT=/workspace/out/libfuzzer-undefined-x86_64",
+      "SANITIZER=undefined"
+    ],
+    "args": [
+      "bash",
+      "-c",
+      "targets_list > /workspace/targets.list.undefined"
+    ]
+  }
+]
diff --git a/infra/build/functions/trial_build.py b/infra/build/functions/trial_build.py
index 2ea70058..03054ece 100644
--- a/infra/build/functions/trial_build.py
+++ b/infra/build/functions/trial_build.py
@@ -116,6 +116,10 @@ def get_args(args=None):
                       required=False,
                       default=None,
                       help='Use specified OSS-Fuzz branch.')
+  parser.add_argument('--force-build',
+                      action='store_true',
+                      help='Build projects that failed to build on OSS-Fuzz\'s '
+                      'production builder.')
   parsed_args = parser.parse_args(args)
   if 'all' in parsed_args.projects:  # Explicit opt-in for all.
     parsed_args.projects = get_all_projects()
@@ -132,7 +136,7 @@ def get_all_projects():
   ])
 
 
-def get_projects_to_build(specified_projects, build_type):
+def get_projects_to_build(specified_projects, build_type, force_build):
   """Returns the list of projects that should be built based on the projects
   specified by the user (|specified_projects|) the |project_statuses| of the
   last builds and the |build_type|."""
@@ -140,10 +144,10 @@ def get_projects_to_build(specified_projects, build_type):
 
   project_statuses = _get_production_build_statuses(build_type)
   for project in specified_projects:
-    if project not in project_statuses:
-      buildable_projects.append(project)
-      continue
-    if project_statuses[project]:
+    if (project not in project_statuses or project_statuses[project] or
+        force_build):
+      # If we don't have data on the project, then we have no reason not to
+      # build it.
       buildable_projects.append(project)
       continue
 
@@ -152,7 +156,7 @@ def get_projects_to_build(specified_projects, build_type):
   return buildable_projects
 
 
-def _do_builds(args, config, credentials, build_type, projects):
+def _do_build_type_builds(args, config, credentials, build_type, projects):
   """Does |build_type| test builds of |projects|."""
   build_ids = {}
   for project_name in projects:
@@ -214,6 +218,7 @@ def check_finished(build_id, project, cloudbuild_api, cloud_project,
   build_status = get_build_status_from_gcb(cloudbuild_api, cloud_project,
                                            build_id)
   if build_status not in FINISHED_BUILD_STATUSES:
+    logging.debug('build: %d not finished.', build_id)
     return False
   build_results[project] = build_status == 'SUCCESS'
   return True
@@ -231,13 +236,17 @@ def wait_on_builds(build_ids, credentials, cloud_project):
   wait_builds = build_ids.copy()
   build_results = {}
   while wait_builds:
-    logging.info('Polling')
-    for project, build_id in list(wait_builds.items()):
-      if check_finished(build_id, project, cloudbuild_api, cloud_project,
-                        build_results):
-        del wait_builds[project]
-      time.sleep(1)  # Avoid rate limiting.
-    print(wait_builds)
+    logging.info('Polling: %s', wait_builds)
+    for project, project_build_ids in list(wait_builds.items()):
+      for build_id in project_build_ids[:]:
+        if check_finished(build_id, project, cloudbuild_api, cloud_project,
+                          build_results):
+
+          wait_builds[project].remove(build_id)
+          if not wait_builds[project]:
+            del wait_builds[project]
+
+        time.sleep(1)  # Avoid rate limiting.
 
   print('Printing results')
   print('Project, Statuses')
@@ -247,7 +256,7 @@ def wait_on_builds(build_ids, credentials, cloud_project):
   return all(build_results.values())
 
 
-def do_test_builds(args):
+def _do_test_builds(args):
   """Does test coverage and fuzzing builds."""
   # TODO(metzman): Make this handle concurrent builds.
   build_types = []
@@ -260,8 +269,10 @@ def do_test_builds(args):
     build_types.append(BUILD_TYPES['introspector'])
   if sanitizers:
     build_types.append(BUILD_TYPES['fuzzing'])
+  build_ids = collections.defaultdict(list)
   for build_type in build_types:
-    projects = get_projects_to_build(list(args.projects), build_type)
+    projects = get_projects_to_build(list(args.projects), build_type,
+                                     args.force_build)
     config = build_project.Config(testing=True,
                                   test_image_suffix=TEST_IMAGE_SUFFIX,
                                   branch=args.branch,
@@ -269,7 +280,11 @@ def do_test_builds(args):
                                   upload=False)
     credentials = (
         oauth2client.client.GoogleCredentials.get_application_default())
-    build_ids = _do_builds(args, config, credentials, build_type, projects)
+    project_builds = _do_build_type_builds(args, config, credentials,
+                                           build_type, projects)
+    for project, project_build_id in project_builds.items():
+      build_ids[project].append(project_build_id)
+
   return wait_on_builds(build_ids, credentials, IMAGE_PROJECT)
 
 
@@ -282,7 +297,7 @@ def trial_build_main(args=None, local_base_build=True):
         TEST_IMAGE_SUFFIX)
   else:
     build_and_push_test_images.gcb_build_and_push_images(TEST_IMAGE_SUFFIX)
-  return do_test_builds(args)
+  return _do_test_builds(args)
 
 
 def main():
diff --git a/infra/build/functions/trial_build/cloudbuild.yaml b/infra/build/functions/trial_build/cloudbuild.yaml
index 8f7122d4..5a5a7d04 100644
--- a/infra/build/functions/trial_build/cloudbuild.yaml
+++ b/infra/build/functions/trial_build/cloudbuild.yaml
@@ -11,6 +11,7 @@ steps:
   args: []
   env:
     - 'PULL_REQUEST_NUMBER=${_PR_NUMBER}'
+    - 'BRANCH=${_HEAD_BRANCH}'
   timeout: 21600s # 6 hours
 options:
   logging: CLOUD_LOGGING_ONLY
diff --git a/infra/build/functions/trial_build_test.py b/infra/build/functions/trial_build_test.py
new file mode 100644
index 00000000..0827dc74
--- /dev/null
+++ b/infra/build/functions/trial_build_test.py
@@ -0,0 +1,86 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+################################################################################
+"""Tests for trial_build.py."""
+import json
+import unittest
+from unittest import mock
+
+import test_utils
+import trial_build
+
+
+class GetProjectsToBuild(unittest.TestCase):
+  """Tests for get_projects_to_build."""
+
+  PROJECTS = ['myproject', 'myfailingproject']
+
+  @mock.patch('trial_build._get_production_build_statuses',
+              return_value={
+                  'myproject': True,
+                  'myfailingproject': False
+              })
+  def test_force_build(self, mock_get_production_build_statuses):
+    """Tests force build works."""
+    del mock_get_production_build_statuses
+    buildable_projects = trial_build.get_projects_to_build(
+        self.PROJECTS, 'fuzzing', True)
+    self.assertEqual(self.PROJECTS, buildable_projects)
+
+  @mock.patch('trial_build._get_production_build_statuses',
+              return_value={
+                  'myproject': True,
+                  'myfailingproject': False
+              })
+  def test_get_projects_to_build(self, mock_get_production_build_statuses):
+    """Tests get_projects_to_build works."""
+    del mock_get_production_build_statuses
+    buildable_projects = trial_build.get_projects_to_build(
+        self.PROJECTS, 'fuzzing', True)
+    self.assertEqual(self.PROJECTS, buildable_projects)
+
+
+class TrialBuildMainTest(unittest.TestCase):
+  """Tests for trial_build_main."""
+
+  @mock.patch('trial_build.wait_on_builds', return_value=True)
+  @mock.patch('oauth2client.client.GoogleCredentials.get_application_default',
+              return_value=None)
+  @mock.patch('build_project.run_build')
+  @mock.patch('build_and_push_test_images.build_and_push_images')
+  def test_build_steps_correct(self, mock_gcb_build_and_push_images,
+                               mock_run_build, mock_get_application_default,
+                               mock_wait_on_builds):
+    """Tests that the correct build steps for building a project are passed to
+    GCB."""
+    del mock_gcb_build_and_push_images
+    del mock_get_application_default
+    del mock_wait_on_builds
+    self.maxDiff = None  # pylint: disable=invalid-name
+    build_id = 1
+    mock_run_build.return_value = build_id
+    branch_name = 'mybranch'
+    project = 'skcms'
+    args = [
+        '--sanitizers', 'address', 'undefined', '--fuzzing-engines', 'afl',
+        'libfuzzer', '--branch', branch_name, '--force-build', project
+    ]
+    self.assertTrue(trial_build.trial_build_main(args))
+    expected_build_steps_path = test_utils.get_test_data_file_path(
+        'expected_trial_build_steps.json')
+    with open(expected_build_steps_path, 'r') as file_handle:
+      expected_build_steps = json.load(file_handle)
+    self.assertEqual(mock_run_build.call_args_list[0][0][1],
+                     expected_build_steps)