Add comments and other minor refactors.

Author: thompsonja

test/cluster/keytar/keytar.py

@@ -1,5 +1,10 @@
 #!/usr/bin/env python
-"""Keytar flask app."""
+"""Keytar flask app.
+
+This program is responsible for exposing an interface to trigger cluster level
+tests. For instance, docker webhooks can be configured to point to this
+application in order to trigger tests upon pushing new docker images.
+"""
 
 import argparse
 import datetime
@@ -22,12 +27,16 @@
 results = {}
 
 
+class KeytarError(Exception):
+  pass
+
+
 def run_test_config(config):
   """Runs a single test iteration from a configuration."""
   tempdir = tempfile.mkdtemp()
   logging.info('Fetching github repository')
 
-  # Get the github repo
+  # Get the github repo and clone it.
   github_config = config['github']
   github_clone_args, github_repo_dir = _get_download_github_repo_args(
       tempdir, github_config)
@@ -40,7 +49,8 @@ def run_test_config(config):
   _add_new_result(timestamp)
   results[timestamp]['docker_image'] = config['docker_image']
 
-  # Run initialization
+  # Generate a test script with the steps described in the configuration,
+  # as well as the command to execute the test_runner.
   with tempfile.NamedTemporaryFile(dir=tempdir, delete=False) as f:
     tempscript = f.name
     f.write('#!/bin/bash\n')
@@ -71,6 +81,7 @@ def index():
 
 @app.route('/test_results')
 def test_results():
+  # Get all test results or a specific test result in JSON.
   if 'name' not in flask.request.args:
     return json.dumps([results[x] for x in sorted(results)])
   else:
@@ -80,13 +91,15 @@ def test_results():
 
 @app.route('/test_log')
 def test_log():
+  # Fetch the output from a test.
   log = '%s.log' % os.path.basename(flask.request.values['log_name'])
   return (flask.send_from_directory('/tmp/testlogs', log), 200,
           {'Content-Type': 'text/css'})
 
 
 @app.route('/update_results', methods=['POST'])
 def update_results():
+  # Update the results dict, called from the test_runner.
   update_args = flask.request.get_json()
   time = update_args['time']
   for k, v in update_args.iteritems():
@@ -119,59 +132,74 @@ def test_request():
     flask.abort(400, validation_error)
   webhook_data = flask.request.get_json()
   repo_name = webhook_data['repository']['repo_name']
-  configs = [c for c in keytar_config if c['docker_image'] == repo_name]
-  if not configs:
+  test_configs = [
+      c for c in keytar_config['config'] if c['docker_image'] == repo_name]
+  if not test_configs:
     return 'No config found for repo_name: %s' % repo_name
-  for config in configs:
-    test_worker.add_test(config)
+  for test_config in test_configs:
+    test_worker.add_test(test_config)
   return 'OK'
 
 
-def process_config(config):
-  global keytar_config
-  keytar_config = config['config']
-  if 'install' in config:
-    install_config = config['install']
-    if 'cluster_setup' in install_config:
-      for cluster_setup in install_config['cluster_setup']:
-        if cluster_setup['type'] == 'gke':
-          gcloud_args = ['gcloud', 'auth', 'activate-service-account',
-                         '--key-file', cluster_setup['keyfile']]
-          logging.info('authenticating using keyfile: %s',
-                       cluster_setup['keyfile'])
-          subprocess.call(gcloud_args)
-          if 'project_name' in cluster_setup:
-            logging.info('Setting gcloud project to %s',
-                         cluster_setup['project_name'])
-            subprocess.call(
-                ['gcloud', 'config', 'set', 'project',
-                 cluster_setup['project_name']])
-          else:
-            config = subprocess.check_output(
-                ['gcloud', 'config', 'list', '--format', 'json'])
-            project_name = json.loads(config)['core']['project']
-            if not project_name:
-              projects = subprocess.check_output(
-                  ['gcloud', 'projects', 'list'])
-              first_project = projects[0]['projectId']
-              logging.info('gcloud project is unset, setting it to %s',
-                           first_project)
-              subprocess.check_output(
-                  ['gcloud', 'config', 'set', 'project', first_project])
-          os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = (
-              cluster_setup['keyfile'])
-    if 'dependencies' in install_config:
-      subprocess.call(['apt-get', 'update'])
-      os.environ['DEBIAN_FRONTEND'] = 'noninteractive'
-      for dep in install_config['dependencies']:
-        subprocess.call(
-            ['apt-get', 'install', '-y', '--no-install-recommends', dep])
-    if 'extra' in install_config:
-      for step in install_config['extra']:
-        os.system(step)
-    if 'path' in install_config:
-      for path in install_config['path']:
-        os.environ['PATH'] = '%s:%s' % (path, os.environ['PATH'])
+def handle_install_steps():
+  """Runs all config installation/setup steps."""
+  if 'install' not in keytar_config:
+    return
+  install_config = keytar_config['install']
+  if 'cluster_setup' in install_config:
+    # Handle any cluster setup steps, currently only GKE is supported.
+    for cluster_setup in install_config['cluster_setup']:
+      if cluster_setup['type'] == 'gke':
+        if 'keyfile' not in cluster_setup:
+          raise KeytarError('No keyfile found in GKE cluster setup!')
+        # Add authentication steps to allow keytar to start clusters on GKE.
+        gcloud_args = ['gcloud', 'auth', 'activate-service-account',
+                       '--key-file', cluster_setup['keyfile']]
+        logging.info('authenticating using keyfile: %s',
+                     cluster_setup['keyfile'])
+        subprocess.call(gcloud_args)
+        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = (
+            cluster_setup['keyfile'])
+
+        # Ensure that a project name is correctly set. Use the name if provided
+        # in the configuration, otherwise use the current project name, or else
+        # the first available project name.
+        if 'project_name' in cluster_setup:
+          logging.info('Setting gcloud project to %s',
+                       cluster_setup['project_name'])
+          subprocess.call(
+              ['gcloud', 'config', 'set', 'project',
+               cluster_setup['project_name']])
+        else:
+          config = subprocess.check_output(
+              ['gcloud', 'config', 'list', '--format', 'json'])
+          project_name = json.loads(config)['core']['project']
+          if not project_name:
+            projects = subprocess.check_output(
+                ['gcloud', 'projects', 'list'])
+            first_project = projects[0]['projectId']
+            logging.info('gcloud project is unset, setting it to %s',
+                         first_project)
+            subprocess.check_output(
+                ['gcloud', 'config', 'set', 'project', first_project])
+
+  # Install any dependencies using apt-get.
+  if 'dependencies' in install_config:
+    subprocess.call(['apt-get', 'update'])
+    os.environ['DEBIAN_FRONTEND'] = 'noninteractive'
+    for dep in install_config['dependencies']:
+      subprocess.call(
+          ['apt-get', 'install', '-y', '--no-install-recommends', dep])
+
+  # Run any additional commands if provided.
+  if 'extra' in install_config:
+    for step in install_config['extra']:
+      os.system(step)
+
+  # Update path environment variable.
+  if 'path' in install_config:
+    for path in install_config['path']:
+      os.environ['PATH'] = '%s:%s' % (path, os.environ['PATH'])
 
 
 def _add_new_result(timestamp):
@@ -180,6 +208,16 @@ def _add_new_result(timestamp):
 
 
 def _get_download_github_repo_args(tempdir, github_config):
+  """Get arguments for github actions.
+
+  Args:
+    tempdir: Base directory to git clone into.
+    github_config: Configuration describing the repo, branches, etc.
+
+  Returns:
+    ([string], string) for arguments to pass to git, and the directory to
+    clone into.
+  """
   repo_prefix = 'github'
   if 'repo_prefix' in github_config:
     repo_prefix = github_config['repo_prefix']
@@ -192,9 +230,9 @@ def _get_download_github_repo_args(tempdir, github_config):
 
 
 class TestWorker(object):
+  """A simple test queue. HTTP requests append to this work queue."""
 
   def __init__(self):
-    # Create a simple test queue. HTTP requests simply append to the queue.
     self.test_queue = Queue.Queue()
     self.worker_thread = threading.Thread(target=self.worker_loop)
     self.worker_thread.daemon = True
@@ -226,7 +264,8 @@ def add_test(self, config):
     yaml_config = yaml_file.read()
   if not yaml_config:
     raise ValueError('No valid yaml config!')
-  process_config(yaml.load(yaml_config))
+  keytar_config = yaml.load(yaml_config)
+  handle_install_steps()
 
   if not os.path.isdir('/tmp/testlogs'):
     os.mkdir('/tmp/testlogs')

test/cluster/keytar/keytar_web_test.py

@@ -43,7 +43,7 @@ def setUpClass(cls):
     cls.driver = create_webdriver()
     cls.flask_process = subprocess.Popen(
         ['./keytar.py', '--config_file=test_config.yaml',
-         '--port=8080', '--key=foo'],
+         '--port=8080', '--password=foo'],
         preexec_fn=os.setsid)
 
   @classmethod
@@ -63,7 +63,7 @@ def _wait_for_complete_status(self, timeout_s=180):
 
   def test_keytar_web(self):
     self.driver.get('http://localhost:8080')
-    req = urllib2.Request('http://localhost:8080/test_request?key=foo')
+    req = urllib2.Request('http://localhost:8080/test_request?password=foo')
     req.add_header('Content-Type', 'application/json')
     urllib2.urlopen(
         req, json.dumps({'repository': {'repo_name': 'test/image'}}))

test/cluster/keytar/static/script.js

@@ -29,7 +29,7 @@ $(document).ready(function() {
     });
     html += "</tbody></table>";
     resultsElement.append(html);
-  }
+  };
 
   // Poll every second.
   var fetchTestResults = function() {

test/cluster/keytar/test_runner.py

@@ -1,4 +1,10 @@
-"""Main python file."""
+"""Script to run a single cluster test.
+
+This includes the following steps:
+  1. Starting a test cluster (GKE supported).
+  2. Running tests against the cluster.
+  3. Reporting test results.
+"""
 
 import argparse
 import json
@@ -13,46 +19,56 @@
 
 
 def update_result(k, v):
+  """Post a key/value pair test result update."""
   url = '%s/update_results' % keytar_args.server
   req = urllib2.Request(url)
   req.add_header('Content-Type', 'application/json')
   urllib2.urlopen(req, json.dumps({k: v, 'time': keytar_args.timestamp}))
 
 
-def start_sandbox(environment_config, name):
+def run_sandbox_action(environment_config, name, action):
+  """Run a sandbox action (Start/Stop).
+
+  Args:
+    environment_config: yaml configuration for the sandbox.
+    name: unique name for the sandbox.
+    action: action to pass to the sandbox action parameter.
+  """
+  if 'sandbox' not in environment_config:
+    return
+  # Execute sandbox command
+  sandbox_file = os.path.join(repo_dir, environment_config['sandbox'])
+  os.chdir(os.path.dirname(sandbox_file))
+  sandbox_args = [
+      './%s' % os.path.basename(sandbox_file),
+      '-e', environment_config['cluster_type'], '-n', name, '-k', name,
+      '-c', os.path.join(repo_dir, environment_config['config']),
+      '-a', action]
+  update_result('status', 'Running sandbox action: %s' % action)
   try:
-    # Start a sandbox
-    sandbox_file = os.path.join(repo_dir, environment_config['sandbox'])
-    os.chdir(os.path.dirname(sandbox_file))
-    sandbox_args = [
-        './%s' % os.path.basename(sandbox_file),
-        '-e', environment_config['cluster_type'], '-n', name, '-k', name,
-        '-c', os.path.join(repo_dir, environment_config['config'])]
-    update_result('status', 'Starting Sandbox')
-    try:
-      subprocess.check_call(sandbox_args + ['-a', 'Start'])
-    except subprocess.CalledProcessError as e:
-      logging.info('Failed to start sandbox: %s', e.output)
-      update_result('status', 'Sandbox failure')
-
-    logging.info('Running tests')
-    update_result('status', 'Running Tests')
-  except Exception as e:  # pylint: disable=broad-except
-    logging.info('Exception caught: %s', str(e))
-    update_result('status', 'System Error')
+    subprocess.check_call(sandbox_args)
+  except subprocess.CalledProcessError as e:
+    logging.info('Failed to run sandbox action %s: %s', (action, e.output))
+    update_result('status', 'Sandbox failure')
 
 
 def run_test_config():
-  """Runs a single test iteration from a configuration."""
+  """Runs a single test iteration from a configuration.
+
+  This includes bringing up an environment, running the tests, and reporting
+  status.
+  """
+  # Generate a random name. Kubernetes/GKE has name length limits.
   name = 'keytar%s' % format(uuid.uuid4().fields[0], 'x')
   update_result('name', name)
 
   environment_config = config['environment']
-  if 'sandbox' in environment_config:
-    start_sandbox(environment_config, name)
+  run_sandbox_action(environment_config, name, 'Start')
+  logging.info('Running tests')
+  update_result('status', 'Running Tests')
 
   try:
-    # Run tests
+    # Run tests and update results.
     test_results = {}
     for test in config['tests']:
       test_file = os.path.join(repo_dir, test['file'])
@@ -78,24 +94,15 @@ def run_test_config():
     update_result('status', 'Tests Complete')
   except Exception as e:  # pylint: disable=broad-except
     logging.info('Exception caught: %s', str(e))
-    update_result('status', 'System Error')
+    update_result('status', 'System Error running tests: %s' % str(e))
   finally:
-    if 'sandbox' in environment_config:
-      sandbox_file = os.path.join(repo_dir, environment_config['sandbox'])
-      sandbox_args = [
-          './%s' % os.path.basename(sandbox_file),
-          '-e', environment_config['cluster_type'], '-n', name, '-k', name,
-          '-c', os.path.join(repo_dir, environment_config['config'])]
-      os.chdir(os.path.dirname(sandbox_file))
-      update_result('status', 'Teardown')
-      subprocess.call(sandbox_args + ['-a', 'Stop'])
-    update_result('status', 'Complete')
+    run_sandbox_action(environment_config, name, 'Stop')
 
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)
   parser = argparse.ArgumentParser(description='Run keytar')
-  parser.add_argument('-c', '--config', help='Keytar config json')
+  parser.add_argument('-c', '--config', help='Keytar config yaml')
   parser.add_argument('-t', '--timestamp', help='Timestamp string')
   parser.add_argument('-d', '--dir', help='temp dir created for the test')
   parser.add_argument('-s', '--server', help='keytar server address')