Modify the sharding supervisor to support sharding only a sub-shard of a test suite.

sharding_supervisor_unittest.py

+#!/usr/bin/env python
+# Copyright (c) 2012 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Verify basic usage of sharding_supervisor."""
+
+import os
+import subprocess
+import sys
+import unittest
+
+import sharding_supervisor
+
+SHARDING_SUPERVISOR = os.path.join(os.path.dirname(sys.argv[0]),
+                                   'sharding_supervisor.py')
+DUMMY_TEST = os.path.join(os.path.dirname(sys.argv[0]), 'dummy_test.py')
+NUM_CORES = sharding_supervisor.DetectNumCores()
+SHARDS_PER_CORE = sharding_supervisor.SS_DEFAULT_SHARDS_PER_CORE
+
+
+def generate_expected_output(start, end, num_shards):
+  """Generate the expected stdout and stderr for the dummy test."""
+  stdout = ''
+  for i in range(start, end):
+    stdout += 'Running shard %d of %d\n' % (i, num_shards)
+  stderr = '\nALL SHARDS PASSED!\nALL TESTS PASSED!\n'
+
+  return (stdout, stderr)
+
+
+class ShardingSupervisorUnittest(unittest.TestCase):
+  def test_basic_run(self):
+    # Default test.
+    expected_shards = NUM_CORES * SHARDS_PER_CORE
+    (expected_out, expected_err) = \
+        generate_expected_output(0, expected_shards, expected_shards)
+    p = subprocess.Popen([SHARDING_SUPERVISOR, '--no-color', DUMMY_TEST],
+                         stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+    (out, err) = p.communicate()
+    self.assertEqual(expected_out, out)
+    self.assertEqual(expected_err, err)
+    self.assertEqual(0, p.returncode)
+
+  def test_shard_per_core(self):
+    """Test the --shards_per_core parameter."""
+    expected_shards = NUM_CORES * 25
+    (expected_out, expected_err) = \
+        generate_expected_output(0, expected_shards, expected_shards)
+    p = subprocess.Popen([SHARDING_SUPERVISOR, '--no-color',
+                          '--shards_per_core', '25', DUMMY_TEST],
+                         stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+    (out, err) = p.communicate()
+    self.assertEqual(expected_out, out)
+    self.assertEqual(expected_err, err)
+    self.assertEqual(0, p.returncode)
+
+  def test_slave_sharding(self):
+    """Test the --total-slaves and --slave-index parameters."""
+    total_shards = 6
+    expected_shards = NUM_CORES * SHARDS_PER_CORE * total_shards
+
+    # Test every single index to make sure they run correctly.
+    for index in range(total_shards):
+      begin = NUM_CORES * SHARDS_PER_CORE * index
+      end = begin + NUM_CORES * SHARDS_PER_CORE
+      (expected_out, expected_err) = \
+          generate_expected_output(begin, end, expected_shards)
+      p = subprocess.Popen([SHARDING_SUPERVISOR, '--no-color',
+                            '--total-slaves', str(total_shards),
+                            '--slave-index', str(index),
+                            DUMMY_TEST],
+                           stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+      (out, err) = p.communicate()
+      self.assertEqual(expected_out, out)
+      self.assertEqual(expected_err, err)
+      self.assertEqual(0, p.returncode)
+
+
+if __name__ == '__main__':
+  unittest.main()