Add sharding param to pyauto tests

For example,
If --shards=0/4 is specified, divide the list of tests into 4 parts
and use the 1st part.

This is so that tests can be distributed on several machines, by using the right sharding params.

BUG=None
TEST=None

R=dennisjeffrey@chromium.org

Review URL: https://chromiumcodereview.appspot.com/9452006

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@124088 0039d316-1c4b-4281-b951-d872f2087c98

2 files changed, 37 insertions, 1 deletions

diff --git a/chrome/test/pyautolib/pyauto.py b/chrome/test/pyautolib/pyauto.py
index 72d1372..959ca45 100755
--- a/chrome/test/pyautolib/pyauto.py
+++ b/chrome/test/pyautolib/pyauto.py
@@ -33,6 +33,7 @@ import optparse
 import os
 import pickle
 import pprint
+import re
 import shutil
 import signal
 import socket
@@ -4778,6 +4779,10 @@ class Main(object):
         '-L', '--list-tests', action='store_true', default=False,
         help='List all tests, and exit.')
     parser.add_option(
+        '--shard',
+        help='Specify sharding params. Example: 1/3 implies split the list of '
+             'tests into 3 groups of which this is the 1st.')
+    parser.add_option(
         '', '--log-file', type='string', default=None,
         help='Provide a path to a file to which the logger will log')
     parser.add_option(
@@ -5009,8 +5014,21 @@ class Main(object):
     # Set CHROME_HEADLESS. It enables crash reporter on posix.
     os.environ['CHROME_HEADLESS'] = '1'
     os.environ['EXTRA_CHROME_FLAGS'] = chrome_flags
-    pyauto_suite = PyUITestSuite(suite_args)
     test_names = self._ExpandTestNames(self._args)
+
+    # Shard, if requested (--shard).
+    if self._options.shard:
+      matched = re.match('(\d+)/(\d+)', self._options.shard)
+      if not matched:
+        print >>sys.stderr, 'Invalid sharding params: %s' % self._options.shard
+        sys.exit(1)
+      shard_index = int(matched.group(1)) - 1
+      num_shards = int(matched.group(2))
+      if shard_index < 0 or shard_index >= num_shards:
+        print >>sys.stderr, 'Invalid sharding params: %s' % self._options.shard
+        sys.exit(1)
+      test_names = pyauto_utils.Shard(test_names, shard_index, num_shards)
+
     test_names *= self._options.repeat
     logging.debug("Loading %d tests from %s", len(test_names), test_names)
     if self._options.list_tests:  # List tests and exit
@@ -5018,6 +5036,7 @@ class Main(object):
         print name
       sys.exit(0)
     loaded_tests = unittest.defaultTestLoader.loadTestsFromNames(test_names)
+    pyauto_suite = PyUITestSuite(suite_args)
     pyauto_suite.addTests(loaded_tests)
     verbosity = 1
     if self._options.verbose:
diff --git a/chrome/test/pyautolib/pyauto_utils.py b/chrome/test/pyautolib/pyauto_utils.py
index e62ac75..159069a 100644
--- a/chrome/test/pyautolib/pyauto_utils.py
+++ b/chrome/test/pyautolib/pyauto_utils.py
@@ -159,3 +159,20 @@ def PrintPerfResult(graph_name, series_name, data_point, units,
       waterfall_indicator, graph_name, series_name,
       str(data_point).replace(' ', ''), units)
   sys.stdout.flush()
+
+
+def Shard(ilist, shard_index, num_shards):
+  """Shard a given list and return the group at index |shard_index|.
+
+  Args:
+    ilist: input list
+    shard_index: 0-based sharding index
+    num_shards: shard count
+  """
+  chunk_size = len(ilist) / num_shards
+  chunk_start = shard_index * chunk_size
+  if shard_index == num_shards - 1:  # Exhaust the remainder in the last shard.
+    chunk_end = len(ilist)
+  else:
+    chunk_end = chunk_start + chunk_size
+  return ilist[chunk_start:chunk_end]