lib/launch_control: Add library.

Move two libraries for working with Launch Control over, so that they
can be shared between multiple Site Specific cbuildbot builders.

BUG=chromium:547996
TEST=Unitests

Change-Id: I812c7573494dfed149d2cc6ffcb184abb78a6f60
Reviewed-on: https://chromium-review.googlesource.com/308754
Commit-Ready: Don Garrett <dgarrett@chromium.org>
Tested-by: Don Garrett <dgarrett@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>

6 files changed, 312 insertions, 0 deletions

diff --git a/lib/launch_control/__init__.py b/lib/launch_control/__init__.py
new file mode 100644
index 000000000..e69de29bb
--- /dev/null
+++ b/lib/launch_control/__init__.py
diff --git a/lib/launch_control/launch_control b/lib/launch_control/launch_control
new file mode 120000
index 000000000..ef3e37b67
--- /dev/null
+++ b/lib/launch_control/launch_control
@@ -0,0 +1 @@
+../../scripts/wrapper.py
\ No newline at end of file
diff --git a/lib/launch_control/launch_control.py b/lib/launch_control/launch_control.py
new file mode 100644
index 000000000..da829cd4f
--- /dev/null
+++ b/lib/launch_control/launch_control.py
@@ -0,0 +1,142 @@
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Select an Android build, and download symbols for it."""
+
+from __future__ import print_function
+
+import json
+import apiclient
+
+from httplib2 import Http
+from apiclient.discovery import build
+from oauth2client.client import SignedJwtAssertionCredentials
+
+from chromite.lib import commandline
+from chromite.lib import cros_logging as logging
+
+
+def OpenBuildApiProxy(json_key_file):
+  """Open an Android Internal Build API Apiary proxy.
+
+  Will NOT error out if authentication fails until the first real request is
+  made.
+
+  Args:
+    json_key_file: A Json key file to authenticate with. Retrieved from the
+                   Google Developer Console associated with the account to
+                   be used for the requests.
+
+  Returns:
+    Proxy object used to make requests against the API.
+  """
+  # Load the private key associated with the Google service account.
+  with open(json_key_file) as json_file:
+    json_data = json.load(json_file)
+    credentials = SignedJwtAssertionCredentials(
+        json_data['client_email'],
+        json_data['private_key'],
+        'https://www.googleapis.com/auth/androidbuild.internal')
+
+  # Open an authorized API proxy.
+  # See https://g3doc.corp.google.com/wireless/android/build_tools/
+  #         g3doc/public/build_data.md
+  http_auth = credentials.authorize(Http())
+  return build('androidbuildinternal', 'v2beta1', http=http_auth)
+
+
+def FindRecentBuildIds(build_api_proxy, branch, target):
+  """Fetch a list of successful completed build ids for a given branch/target.
+
+  This roughly matches the contents of the first page of build results on the
+  launch control website, except filtered for only successful/completed builds.
+
+  Since builds sometimes complete out of order, new builds can be added to the
+  list out of order.
+
+  Args:
+    build_api_proxy: Result of a previous call to OpenBuildApiProxy.
+    branch: Name of branch to search. Ex. 'git_mnc-dr-ryu-release'
+    target: Build target to search. Ex. 'ryu-userdebug'
+
+  Returns:
+    List of build_ids as integers.
+  """
+  result = build_api_proxy.build().list(
+      buildType='submitted',
+      branch=branch,
+      buildAttemptStatus='complete',
+      successful=True,
+      target=target,
+  ).execute()
+
+  # Extract the build_ids, arrange oldest to newest.
+  return sorted(int(b['buildId']) for b in result['builds'])
+
+
+def FetchBuildArtifact(build_api_proxy, build_id, target, resource_id,
+                       output_file):
+  """Fetch debug symbols associated with a given build.
+
+  Args:
+    build_api_proxy: Result of a previous call to OpenBuildApiProxy.
+    build_id: id of the build to fetch symbols for.
+    target: Build to target fetch symbols for. Ex. 'ryu-userdebug'
+    resource_id: Resource id to fetch. Ex. 'ryu-symbols-2282124.zip'
+    output_file: Path to where to write out the downloaded artifact.
+  """
+  # Open the download connection.
+  download_req = build_api_proxy.buildartifact().get_media(
+      buildId=build_id,
+      target=target,
+      attemptId='latest',
+      resourceId=resource_id)
+
+  # Download the symbols file contents.
+  with open(output_file, mode='wb') as fh:
+    downloader = apiclient.http.MediaIoBaseDownload(
+        fh, download_req, chunksize=20 * 1024 * 1024)
+    done = False
+    while not done:
+      _status, done = downloader.next_chunk()
+
+
+def main(argv):
+  """Command line wrapper for integration testing of the above library.
+
+  This library requires the ability to authenticate to an external service that
+  is restricted from the general public. So, allow manual integration testing by
+  users that have the necessary credentials.
+  """
+  parser = commandline.ArgumentParser(description=__doc__)
+
+  parser.add_argument('--json-key-file', type='path', required=True,
+                      help='Json key file for authenticating to service.')
+  parser.add_argument('--symbols-file', type='path', default='symbols.zip',
+                      help='Where to write symbols file out.')
+  parser.add_argument('--branch', type=str, default='git_mnc-dr-ryu-release',
+                      help='Branch to locate build for.')
+  parser.add_argument('--target', type=str, default='ryu-userdebug',
+                      help='Target to locate build for.')
+
+  opts = parser.parse_args(argv)
+  opts.Freeze()
+
+  build_proxy = OpenBuildApiProxy(opts.json_key_file)
+
+  build_ids = FindRecentBuildIds(
+      build_proxy,
+      branch=opts.branch,
+      target=opts.target)
+
+  build_id = build_ids[0]
+
+  # 'ryu-userdebug' -> 'ryu'
+  board = opts.target.split('-')[0]
+  # E.g. 'ryu-symbols-2282124.zip'
+  resource_id = '%s-symbols-%s.zip' % (board, build_id)
+
+  logging.info('Selected buildId: %s', build_id)
+  FetchBuildArtifact(build_proxy, build_id, opts.target, resource_id,
+                     opts.symbols_file)
diff --git a/lib/launch_control/processed_builds.py b/lib/launch_control/processed_builds.py
new file mode 100644
index 000000000..3e8f614b9
--- /dev/null
+++ b/lib/launch_control/processed_builds.py
@@ -0,0 +1,84 @@
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Select an Android build, and download symbols for it."""
+
+from __future__ import print_function
+
+import json
+
+from chromite.lib import osutils
+
+
+class ProcessedBuildsStorage(object):
+  """A context manager for storing processed builds.
+
+  This is a context manager that loads recent builds, and allows them to be
+  manipulated, and then saves them on exit. Processed builds are stored per
+  branch/target as a list of integers.
+  """
+  def __init__(self, filename):
+    self.filename = filename
+    self.value = self._read()
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    self._write(self.value)
+
+  def _read(self):
+    """Load from disk, and default to an empty store on error."""
+    try:
+      return json.loads(osutils.ReadFile(self.filename))
+    except (ValueError, IOError):
+      # If there was no file, or it was corrupt json, return default.
+      return {}
+
+  def _write(self, new_value):
+    """Write the current store to disk."""
+    return osutils.WriteFile(self.filename,
+                             json.dumps(new_value, sort_keys=True))
+
+  def GetProcessedBuilds(self, branch, target):
+    """Get a list of builds for a branch/target.
+
+    Args:
+      branch: Name of branch as a string.
+      target: Name of target as a string.
+
+    Returns:
+      List of integers associated with the given branch/target.
+    """
+    self.value.setdefault(branch, {})
+    self.value[branch].setdefault(target, [])
+    return self.value[branch][target]
+
+  def PurgeOldBuilds(self, branch, target, retain_list):
+    """Removes uninteresting builds for a branch/target.
+
+    Any build ids not in the retain list are removed.
+
+    Args:
+      branch: Name of branch as a string.
+      target: Name of target as a string.
+      retain_list: List of build ids that are still relevent.
+    """
+    processed = set(self.GetProcessedBuilds(branch, target))
+    retained_processed = processed.intersection(retain_list)
+    self.value[branch][target] = list(retained_processed)
+
+  def AddProcessedBuild(self, branch, target, build_id):
+    """Adds build_id to list for a branch/target.
+
+    It's safe to add a build_id that is already present.
+
+    Args:
+      branch: Name of branch as a string.
+      target: Name of target as a string.
+      build_id: build_id to add, as an integer.
+    """
+    processed = set(self.GetProcessedBuilds(branch, target))
+    processed.add(build_id)
+    self.value[branch][target] = sorted(processed)
diff --git a/lib/launch_control/processed_builds_unittest b/lib/launch_control/processed_builds_unittest
new file mode 120000
index 000000000..ef3e37b67
--- /dev/null
+++ b/lib/launch_control/processed_builds_unittest
@@ -0,0 +1 @@
+../../scripts/wrapper.py
\ No newline at end of file
diff --git a/lib/launch_control/processed_builds_unittest.py b/lib/launch_control/processed_builds_unittest.py
new file mode 100644
index 000000000..91443b642
--- /dev/null
+++ b/lib/launch_control/processed_builds_unittest.py
@@ -0,0 +1,84 @@
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Test processed_builds."""
+
+from __future__ import print_function
+
+import os
+
+from chromite.lib.launch_control import processed_builds
+from chromite.lib import cros_test_lib
+
+
+# Unitests often need access to internals of the thing they test.
+# pylint: disable=protected-access
+
+class ProcessedBuildsStorageTest(cros_test_lib.TempDirTestCase):
+  """Test our helper library for storing processed build ids."""
+
+  def setUp(self):
+    self.testfile = os.path.join(self.tempdir, 'testfile.json')
+
+  def testStartStop(self):
+    with processed_builds.ProcessedBuildsStorage(self.testfile):
+      pass
+
+    self.assertFileContents(self.testfile, '{}')
+
+  def testFetchEmpty(self):
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      self.assertEqual(ps.GetProcessedBuilds('branch', 'target'), [])
+
+    self.assertFileContents(self.testfile, '{"branch": {"target": []}}')
+
+  def testPurgeEmpty(self):
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      ps.PurgeOldBuilds('branch', 'target', [1, 2, 3])
+
+    self.assertFileContents(self.testfile, '{"branch": {"target": []}}')
+
+  def testAddEmpty(self):
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      ps.AddProcessedBuild('branch', 'target', 1)
+
+    self.assertFileContents(self.testfile, '{"branch": {"target": [1]}}')
+
+  def testMultipleUses(self):
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      ps.AddProcessedBuild('branch', 'target', 1)
+      ps.AddProcessedBuild('branch', 'target', 2)
+
+    self.assertFileContents(self.testfile, '{"branch": {"target": [1, 2]}}')
+
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      # Try adding twice, should only happen once.
+      ps.AddProcessedBuild('branch', 'target', 3)
+      ps.AddProcessedBuild('branch', 'target', 3)
+
+    self.assertFileContents(self.testfile, '{"branch": {"target": [1, 2, 3]}}')
+
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      ps.PurgeOldBuilds('branch', 'target', [2, 3])
+      ps.AddProcessedBuild('branch', 'target', 4)
+
+    self.assertFileContents(self.testfile, '{"branch": {"target": [2, 3, 4]}}')
+
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      self.assertEqual(ps.GetProcessedBuilds('branch', 'target'), [2, 3, 4])
+
+  def testAddMultipleBranchTargets(self):
+    with processed_builds.ProcessedBuildsStorage(self.testfile) as ps:
+      ps.AddProcessedBuild('branch1', 'target', 1)
+      ps.AddProcessedBuild('branch2', 'target', 1)
+      ps.AddProcessedBuild('branch2', 'target', 2)
+      ps.AddProcessedBuild('branch2', 'target2', 3)
+
+      self.assertEqual(ps.GetProcessedBuilds('branch2', 'target'),
+                       [1, 2])
+
+    self.assertFileContents(
+        self.testfile,
+        '{"branch1": {"target": [1]},'
+        ' "branch2": {"target": [1, 2], "target2": [3]}}')