Refactor manage.py to be more readable.

testing/buildbot/manage.py

Index: testing/buildbot/manage.py
diff --git a/testing/buildbot/manage.py b/testing/buildbot/manage.py
index 6d1a874e20afc333bf617672fc5ed552ddebbf0f..98de7ef4b04356b5ff0635db03ee7bdde9a5095a 100755
--- a/testing/buildbot/manage.py
+++ b/testing/buildbot/manage.py
@@ -51,12 +51,8 @@ SKIP = {
 }
 
 
-def upgrade_test(test):
-  """Converts from old style string to new style dict."""
-  if isinstance(test, basestring):
-    return {'test': test}
-  assert isinstance(test, dict)
-  return test
+class Error(Exception):
+  """Processing error."""
 
 
 def get_isolates():
@@ -65,33 +61,159 @@ def get_isolates():
   return [os.path.basename(f) for f in files if f.endswith('.isolate')]
 
 
+def process_builder_convert(data, filename, builder, test_name):
+  """Converts 'test_name' to run on Swarming in 'data'.
+
+  Returns True if 'test_name' was found.
+  """
+  result = False
+  for test in data['gtest_tests']:
+    if test['test'] != test_name:
+      continue
+    test.setdefault('swarming', {})
+    if not test['swarming'].get('can_use_on_swarming_builders'):
+      print('- %s: %s' % (filename, builder))
+      test['swarming']['can_use_on_swarming_builders'] = True
+    result = True
+  return result
+
+
+def process_builder_remaining(data, filename, builder, tests_location):
+  """Calculates tests_location when mode is --remaining."""
+  for test in data['gtest_tests']:
+    name = test['test']
+    if test.get('swarming', {}).get('can_use_on_swarming_builders'):
+      tests_location[name]['count_run_on_swarming'] += 1
+    else:
+      tests_location[name]['count_run_local'] += 1
+      tests_location[name]['local_configs'].setdefault(
+          filename, []).append(builder)
+
+
+def process_file(mode, test_name, tests_location, filepath):
+  """Processes a file.
+
+  The action depends on mode. Updates tests_location.
+
+  Return False if the process exit code should be 1.
+  """
+  filename = os.path.basename(filepath)
+  with open(filepath) as f:
+    content = f.read()
+  try:
+    config = json.loads(content)
+  except ValueError as e:
+    raise Error('Exception raised while checking %s: %s' % (filepath, e))
+
+  for builder, data in sorted(config.iteritems()):
+    if builder in SKIP:
+      # Oddities.
+      continue
+    if not isinstance(data, dict):
+      raise Error('%s: %s is broken: %s' % (filename, builder, data))
+    if 'gtest_tests' not in data:
+      continue
+    if not isinstance(data['gtest_tests'], list):
+      raise Error(
+          '%s: %s is broken: %s' % (filename, builder, data['gtest_tests']))
+    if not all(isinstance(g, dict) for g in data['gtest_tests']):
+      raise Error(
+          '%s: %s is broken: %s' % (filename, builder, data['gtest_tests']))
+
+    config[builder]['gtest_tests'] = sorted(
+        data['gtest_tests'], key=lambda x: x['test'])
+    if mode == 'remaining':
+      process_builder_remaining(data, filename, builder, tests_location)
+    elif mode == 'convert':
+      process_builder_convert(data, filename, builder, test_name)
+
+  expected = json.dumps(
+      config, sort_keys=True, indent=2, separators=(',', ': ')) + '\n'
+  if content != expected:
+    if mode in ('convert', 'write'):
+      with open(filepath, 'wb') as f:
+        f.write(expected)
+      if mode == 'write':
+        print('Updated %s' % filename)
+    else:
+      print('%s is not in canonical format' % filename)
+      print('run `testing/buildbot/manage.py -w` to fix')
+    return mode != 'check'
+  return True
+
+
+def print_remaining(test_name,tests_location):
+  """Prints a visual summary of what tests are yet to be converted to run on
+  Swarming.
+  """
+  if test_name:
+    if test_name not in tests_location:
+      raise Error('Unknown test %s' % test_name)
+    for config, builders in sorted(
+        tests_location[test_name]['local_configs'].iteritems()):
+      print('%s:' % config)
+      for builder in sorted(builders):
+        print('  %s' % builder)
+    return
+
+  isolates = get_isolates()
+  l = max(map(len, tests_location))
+  print('%-*s%sLocal       %sSwarming  %sMissing isolate' %
+      (l, 'Test', colorama.Fore.RED, colorama.Fore.GREEN,
+        colorama.Fore.MAGENTA))
+  total_local = 0
+  total_swarming = 0
+  for name, location in sorted(tests_location.iteritems()):
+    if not location['count_run_on_swarming']:
+      c = colorama.Fore.RED
+    elif location['count_run_local']:
+      c = colorama.Fore.YELLOW
+    else:
+      c = colorama.Fore.GREEN
+    total_local += location['count_run_local']
+    total_swarming += location['count_run_on_swarming']
+    missing_isolate = ''
+    if name + '.isolate' not in isolates:
+      missing_isolate = colorama.Fore.MAGENTA + '*'
+    print('%s%-*s %4d           %4d    %s' %
+        (c, l, name, location['count_run_local'],
+          location['count_run_on_swarming'], missing_isolate))
+
+  total = total_local + total_swarming
+  p_local = 100. * total_local / total
+  p_swarming = 100. * total_swarming / total
+  print('%s%-*s %4d (%4.1f%%)   %4d (%4.1f%%)' %
+      (colorama.Fore.WHITE, l, 'Total:', total_local, p_local,
+        total_swarming, p_swarming))
+  print('%-*s                %4d' % (l, 'Total executions:', total))
+
+
 def main():
   colorama.init()
   parser = argparse.ArgumentParser(description=sys.modules[__name__].__doc__)
   group = parser.add_mutually_exclusive_group(required=True)
   group.add_argument(
-      '-c', '--check', action='store_true', help='Only check the files')
+      '-c', '--check', dest='mode', action='store_const', const='check',
+      default='check', help='Only check the files')
   group.add_argument(
-      '--convert', action='store_true',
+      '--convert', dest='mode', action='store_const', const='convert',
       help='Convert a test to run on Swarming everywhere')
   group.add_argument(
-      '--remaining', action='store_true',
+      '--remaining', dest='mode', action='store_const', const='remaining',
       help='Count the number of tests not yet running on Swarming')
   group.add_argument(
-      '-w', '--write', action='store_true', help='Rewrite the files')
+      '-w', '--write', dest='mode', action='store_const', const='write',
+      help='Rewrite the files')
   parser.add_argument(
       'test_name', nargs='?',
       help='The test name to print which configs to update; only to be used '
            'with --remaining')
   args = parser.parse_args()
 
-  if args.convert or args.remaining:
-    isolates = get_isolates()
-
-  if args.convert:
+  if args.mode == 'convert':
     if not args.test_name:
       parser.error('A test name is required with --convert')
-    if args.test_name + '.isolate' not in isolates:
+    if args.test_name + '.isolate' not in get_isolates():
       parser.error('Create %s.isolate first' % args.test_name)
 
   # Stats when running in --remaining mode;
@@ -100,102 +222,18 @@ def main():
         'count_run_local': 0, 'count_run_on_swarming': 0, 'local_configs': {}
       })
 
-  result = 0
-  for filepath in glob.glob(os.path.join(THIS_DIR, '*.json')):
-    filename = os.path.basename(filepath)
-    with open(filepath) as f:
-      content = f.read()
-    try:
-      config = json.loads(content)
-    except ValueError as e:
-      print "Exception raised while checking %s: %s" % (filepath, e)
-      raise
-    for builder, data in sorted(config.iteritems()):
-      if builder in SKIP:
-        # Oddities.
-        continue
-
-      if not isinstance(data, dict):
-        print('%s: %s is broken: %s' % (filename, builder, data))
-        continue
-
-      if 'gtest_tests' in data:
-        config[builder]['gtest_tests'] = sorted(
-          (upgrade_test(l) for l in data['gtest_tests']),
-          key=lambda x: x['test'])
-
-        if args.remaining:
-          for test in data['gtest_tests']:
-            name = test['test']
-            if test.get('swarming', {}).get('can_use_on_swarming_builders'):
-              tests_location[name]['count_run_on_swarming'] += 1
-            else:
-              tests_location[name]['count_run_local'] += 1
-              tests_location[name]['local_configs'].setdefault(
-                  filename, []).append(builder)
-        elif args.convert:
-          for test in data['gtest_tests']:
-            if test['test'] != args.test_name:
-              continue
-            test.setdefault('swarming', {})
-            if not test['swarming'].get('can_use_on_swarming_builders'):
-              print('- %s: %s' % (filename, builder))
-              test['swarming']['can_use_on_swarming_builders'] = True
-
-    expected = json.dumps(
-        config, sort_keys=True, indent=2, separators=(',', ': ')) + '\n'
-    if content != expected:
-      result = 1
-      if args.write or args.convert:
-        with open(filepath, 'wb') as f:
-          f.write(expected)
-        if args.write:
-          print('Updated %s' % filename)
-      else:
-        print('%s is not in canonical format' % filename)
-        print('run `testing/buildbot/manage.py -w` to fix')
-
-  if args.remaining:
-    if args.test_name:
-      if args.test_name not in tests_location:
-        print('Unknown test %s' % args.test_name)
-        return 1
-      for config, builders in sorted(
-          tests_location[args.test_name]['local_configs'].iteritems()):
-        print('%s:' % config)
-        for builder in sorted(builders):
-          print('  %s' % builder)
-    else:
-      l = max(map(len, tests_location))
-      print('%-*s%sLocal       %sSwarming  %sMissing isolate' %
-          (l, 'Test', colorama.Fore.RED, colorama.Fore.GREEN,
-            colorama.Fore.MAGENTA))
-      total_local = 0
-      total_swarming = 0
-      for name, location in sorted(tests_location.iteritems()):
-        if not location['count_run_on_swarming']:
-          c = colorama.Fore.RED
-        elif location['count_run_local']:
-          c = colorama.Fore.YELLOW
-        else:
-          c = colorama.Fore.GREEN
-        total_local += location['count_run_local']
-        total_swarming += location['count_run_on_swarming']
-        missing_isolate = ''
-        if name + '.isolate' not in isolates:
-          missing_isolate = colorama.Fore.MAGENTA + '*'
-        print('%s%-*s %4d           %4d    %s' %
-            (c, l, name, location['count_run_local'],
-              location['count_run_on_swarming'], missing_isolate))
-
-      total = total_local + total_swarming
-      p_local = 100. * total_local / total
-      p_swarming = 100. * total_swarming / total
-      print('%s%-*s %4d (%4.1f%%)   %4d (%4.1f%%)' %
-          (colorama.Fore.WHITE, l, 'Total:', total_local, p_local,
-            total_swarming, p_swarming))
-      print('%-*s                %4d' % (l, 'Total executions:', total))
-  return result
+  try:
+    result = 0
+    for filepath in glob.glob(os.path.join(THIS_DIR, '*.json')):
+      if not process_file(args.mode, args.test_name, tests_location, filepath):
+        result = 1
+
+    if args.mode == 'remaining':
+      print_remaining(args.test_name, tests_location)
+    return result
+  except Error as e:
+    sys.stderr.write('%s\n' % e)
+    return 1
 
 
 if __name__ == "__main__":