Add support for sharding to AndroidJunitRunner

Buckets tests based on their descriptors'
hashcodes.

Bug: 12929805
Change-Id: Ib4bf43a836a85cef7a1f6f1483eb82514bbf9009

3 files changed, 103 insertions, 0 deletions

diff --git a/support/src/android/support/test/internal/runner/TestRequestBuilder.java b/support/src/android/support/test/internal/runner/TestRequestBuilder.java
index 7cc0a7d..cf0e6c9 100644
--- a/support/src/android/support/test/internal/runner/TestRequestBuilder.java
+++ b/support/src/android/support/test/internal/runner/TestRequestBuilder.java
@@ -195,6 +195,39 @@ public class TestRequestBuilder {
         }
     }
 
+    private static class ShardingFilter extends Filter {
+        private final int mNumShards;
+        private final int mShardIndex;
+
+        ShardingFilter(int numShards, int shardIndex) {
+            mNumShards = numShards;
+            mShardIndex = shardIndex;
+        }
+
+        @Override
+        public boolean shouldRun(Description description) {
+            if (description.isTest()) {
+                return (Math.abs(description.hashCode()) % mNumShards) == mShardIndex;
+            }
+            // this is a suite, explicitly check if any children should run
+            for (Description each : description.getChildren()) {
+                if (shouldRun(each)) {
+                    return true;
+                }
+            }
+            // no children to run, filter this out
+            return false;
+        }
+
+        /**
+         * {@inheritDoc}
+         */
+        @Override
+        public String describe() {
+            return String.format("Shard %s of %s shards", mShardIndex, mNumShards);
+        }
+    }
+
     public TestRequestBuilder(PrintStream writer, String... apkPaths) {
         mApkPaths = apkPaths;
         mTestLoader = new TestLoader(writer);
@@ -305,6 +338,10 @@ public class TestRequestBuilder {
         }
     }
 
+    public void addShardingFilter(int numShards, int shardIndex) {
+        mFilter = mFilter.intersect(new ShardingFilter(numShards, shardIndex));
+    }
+
     /**
      * Build a request that will generate test started and test ended events, but will skip actual
      * test execution.
diff --git a/support/src/android/support/test/runner/AndroidJUnitRunner.java b/support/src/android/support/test/runner/AndroidJUnitRunner.java
index 5a600ed..bffd0c1 100644
--- a/support/src/android/support/test/runner/AndroidJUnitRunner.java
+++ b/support/src/android/support/test/runner/AndroidJUnitRunner.java
@@ -117,6 +117,11 @@ import java.util.List;
  * instrument -w -e notAnnotation com.android.foo.MyAnnotation,com.android.foo.AnotherAnnotation
  * com.android.foo/android.support.test.runner.AndroidJUnitRunner
  * <p/>
+ * <b>Filter test run to a shard of all tests, where numShards is an integer greater than 0 and
+ * shardIndex is an integer between 0 (inclusive) and numShards (exclusive):</b> adb shell am
+ * instrument -w -e numShards 4 -e shardIndex 1
+ * com.android.foo/android.support.test.runner.AndroidJUnitRunner
+ * <p/>
  * <b>To run in 'log only' mode</b>
  * -e log true
  * This option will load and iterate through all test classes and methods, but will bypass actual
@@ -148,6 +153,8 @@ public class AndroidJUnitRunner extends Instrumentation {
     private static final String ARGUMENT_LOG_ONLY = "log";
     private static final String ARGUMENT_ANNOTATION = "annotation";
     private static final String ARGUMENT_NOT_ANNOTATION = "notAnnotation";
+    private static final String ARGUMENT_NUM_SHARDS = "numShards";
+    private static final String ARGUMENT_SHARD_INDEX = "shardIndex";
     private static final String ARGUMENT_DELAY_MSEC = "delay_msec";
     private static final String ARGUMENT_COVERAGE = "coverage";
     private static final String ARGUMENT_COVERAGE_PATH = "coverageFile";
@@ -422,6 +429,23 @@ public class AndroidJUnitRunner extends Instrumentation {
             }
         }
 
+        // Accept either string or int.
+        Object numShardsObj = arguments.get(ARGUMENT_NUM_SHARDS);
+        Object shardIndexObj = arguments.get(ARGUMENT_SHARD_INDEX);
+        if (numShardsObj != null && shardIndexObj != null) {
+            int numShards = -1;
+            int shardIndex = -1;
+            try {
+                numShards = Integer.parseInt(numShardsObj.toString());
+                shardIndex = Integer.parseInt(shardIndexObj.toString());
+            } catch(NumberFormatException e) {
+                Log.e(LOG_TAG, "Invalid sharding parameter", e);
+            }
+            if (numShards > 0 && shardIndex >= 0 && shardIndex < numShards) {
+                builder.addShardingFilter(numShards, shardIndex);
+            }
+        }
+
         if (getBooleanArgument(ARGUMENT_LOG_ONLY)) {
             builder.setSkipExecution(true);
         }
diff --git a/support/tests/src/android/support/test/internal/runner/TestRequestBuilderTest.java b/support/tests/src/android/support/test/internal/runner/TestRequestBuilderTest.java
index 677a437..81cf3f7 100644
--- a/support/tests/src/android/support/test/internal/runner/TestRequestBuilderTest.java
+++ b/support/tests/src/android/support/test/internal/runner/TestRequestBuilderTest.java
@@ -507,4 +507,46 @@ public class TestRequestBuilderTest {
         Assert.assertEquals("testRunThis",
                 result.getFailures().get(0).getDescription().getMethodName());
     }
+
+    /**
+     * Test the sharding filter.
+     */
+    @Test
+    public void testShardingFilter() {
+        JUnitCore testRunner = new JUnitCore();
+
+        TestRequestBuilder[] builders = new TestRequestBuilder[5];
+        Result[] results = new Result[4];
+        int totalRun = 0;
+        // The last iteration through the loop doesn't add a ShardingFilter - it runs all the
+        // tests to establish a baseline for the total number that should have been run.
+        for (int i = 0; i < 5; i++) {
+            TestRequestBuilder b
+                    = new TestRequestBuilder(new PrintStream(new ByteArrayOutputStream()));
+            if (i < 4) {
+                b.addShardingFilter(4, i);
+            }
+            builders[i] = b;
+            b.addTestClass(SampleTest.class.getName());
+            b.addTestClass(SampleNoSize.class.getName());
+            b.addTestClass(SampleClassSize.class.getName());
+            b.addTestClass(SampleJUnit3Test.class.getName());
+            b.addTestClass(SampleOverrideSize.class.getName());
+            b.addTestClass(SampleJUnit3ClassSize.class.getName());
+            b.addTestClass(SampleMultipleAnnotation.class.getName());
+            TestRequest request = b.build(mInstr, mBundle);
+            Result result = testRunner.run(request.getRequest());
+            if (i == 4) {
+                Assert.assertEquals(result.getRunCount(), totalRun);
+            } else {
+                totalRun += result.getRunCount();
+                results[i] = result;
+            }
+        }
+        for (int i = 0; i < 4; i++) {
+            // Theoretically everything could collide into one shard, but, we'll trust that
+            // the implementation of hashCode() is random enough to avoid that.
+            Assert.assertTrue(results[i].getRunCount() < totalRun);
+        }
+    }
 }