Add main_repo to build status.json (google#9822)

Fixes google#9782

infra/build/build_status/update_build_status.py

@@ -25,6 +25,7 @@
 import googleapiclient.errors
 from google.cloud import ndb
 from google.cloud import storage
+import yaml
 
 import build_and_run_coverage
 import build_lib
@@ -203,17 +204,31 @@ def get_build_history(build_ids):
   return project
 
 
+def _get_main_repo(project_name):
+  """Get the main repo for a project."""
+  project = datastore_entities.Project.query(
+      datastore_entities.Project.name == project_name).get()
+  if not project:
+    return None
+
+  project_yaml = yaml.safe_load(project.project_yaml_contents)
+  return project_yaml.get('main_repo')
+
+
 # pylint: disable=too-many-locals
 def update_build_status(build_tag, status_filename):
   """Update build statuses."""
   projects = []
 
   def process_project(project_build):
     """Process a project."""
-    project = get_build_history(project_build.build_ids)
-    project['name'] = project_build.project
-    print('Processing project', project['name'])
-    return project
+    # We need a new context for every thread.
+    with ndb.Client().context():
+      project = get_build_history(project_build.build_ids)
+      project['name'] = project_build.project
+      project['main_repo'] = _get_main_repo(project_build.project)
+      print('Processing project', project['name'])
+      return project
 
   with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
     futures = []

infra/build/build_status/update_build_status_test.py

@@ -38,9 +38,8 @@ class MockGetBuild:
   def __init__(self, builds):
     self.builds = builds
 
-  def get_build(self, cloudbuild, image_project, build_id):
+  def get_build(self, build_id):
     """Mimic build object retrieval."""
-    del cloudbuild, image_project
     for build in self.builds:
       if build['build_id'] == build_id:
         return build
@@ -61,7 +60,7 @@ def test_get_build_history(self, mock_upload_log, mock_cloud_build,
     mock_upload_log.return_value = True
     builds = [{'build_id': '1', 'finishTime': 'test_time', 'status': 'SUCCESS'}]
     mock_get_build = MockGetBuild(builds)
-    update_build_status.get_build = mock_get_build.get_build
+    update_build_status.BuildGetter.get_build = mock_get_build.get_build
 
     expected_projects = {
         'history': [{
@@ -77,25 +76,14 @@ def test_get_build_history(self, mock_upload_log, mock_cloud_build,
     self.assertDictEqual(update_build_status.get_build_history(['1']),
                          expected_projects)
 
-  def test_get_build_history_missing_log(self, mock_upload_log,
-                                         mock_cloud_build, mock_google_auth):
-    """Test for missing build log file."""
-    del mock_cloud_build, mock_google_auth
-    builds = [{'build_id': '1', 'finishTime': 'test_time', 'status': 'SUCCESS'}]
-    mock_get_build = MockGetBuild(builds)
-    update_build_status.get_build = mock_get_build.get_build
-    mock_upload_log.return_value = False
-    self.assertRaises(update_build_status.MissingBuildLogError,
-                      update_build_status.get_build_history, ['1'])
-
   def test_get_build_history_no_last_success(self, mock_upload_log,
                                              mock_cloud_build,
                                              mock_google_auth):
     """Test when there is no last successful build."""
     del mock_cloud_build, mock_google_auth
     builds = [{'build_id': '1', 'finishTime': 'test_time', 'status': 'FAILURE'}]
     mock_get_build = MockGetBuild(builds)
-    update_build_status.get_build = mock_get_build.get_build
+    update_build_status.BuildGetter.get_build = mock_get_build.get_build
     mock_upload_log.return_value = True
 
     expected_projects = {
@@ -236,6 +224,19 @@ def test_update_build_status(self, mock_upload_log, mock_cloud_build,
     mock_upload_log.return_value = True
     status_filename = 'status.json'
     with ndb.Client().context():
+      datastore_entities.Project(
+          name='test-project-1',
+          project_yaml_contents=(
+              'main_repo: "https://github.com/main/repo1"')).put()
+      datastore_entities.Project(
+          name='test-project-2',
+          project_yaml_contents=(
+              'main_repo: "https://github.com/main/repo2"')).put()
+      datastore_entities.Project(
+          name='test-project-3',
+          project_yaml_contents=(
+              'main_repo: "https://github.com/main/repo3"')).put()
+
       datastore_entities.BuildsHistory(id='test-project-1-fuzzing',
                                        build_tag='fuzzing',
                                        project='test-project-1',
@@ -264,7 +265,7 @@ def test_update_build_status(self, mock_upload_log, mock_cloud_build,
           'status': 'WORKING'
       }]
       mock_get_build = MockGetBuild(builds)
-      update_build_status.get_build = mock_get_build.get_build
+      update_build_status.BuildGetter.get_build = mock_get_build.get_build
 
       expected_data = {
           'projects': [{
@@ -273,7 +274,8 @@ def test_update_build_status(self, mock_upload_log, mock_cloud_build,
                   'finish_time': 'test_time',
                   'success': False
               }],
-              'name': 'test-project-2'
+              'name': 'test-project-2',
+              'main_repo': 'https://github.com/main/repo2',
           }, {
               'history': [{
                   'build_id': '1',
@@ -284,10 +286,12 @@ def test_update_build_status(self, mock_upload_log, mock_cloud_build,
                   'build_id': '1',
                   'finish_time': 'test_time'
               },
-              'name': 'test-project-1'
+              'name': 'test-project-1',
+              'main_repo': 'https://github.com/main/repo1',
           }, {
               'history': [],
-              'name': 'test-project-3'
+              'name': 'test-project-3',
+              'main_repo': 'https://github.com/main/repo3',
           }]
       }