From 025f7ae642bbb20fd60eaf094374b56fba5975f1 Mon Sep 17 00:00:00 2001
From: Casey Hillers <caseyhillers@gmail.com>
Date: Tue, 31 Mar 2020 14:01:01 -0700
Subject: [PATCH] Add ability to shard customer tests (#53433)

---
 dev/customer_testing/run_tests.dart | 40 +++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/dev/customer_testing/run_tests.dart b/dev/customer_testing/run_tests.dart
index 2847f27060..6ec914f46b 100644
--- a/dev/customer_testing/run_tests.dart
+++ b/dev/customer_testing/run_tests.dart
@@ -25,6 +25,18 @@ Future<bool> run(List<String> arguments) async {
       help: 'How many times to run each test. Set to a high value to look for flakes.',
       valueHelp: 'count',
     )
+    ..addOption(
+      'shards',
+      defaultsTo: '1',
+      help: 'How many shards to split the tests into. Used in continuous integration.',
+      valueHelp: 'count',
+    )
+    ..addOption(
+      'shard-index',
+      defaultsTo: '0',
+      help: 'The current shard to run the tests with the range [0 .. shards - 1]. Used in continuous integration.',
+      valueHelp: 'count',
+    )
     ..addFlag(
       'skip-on-fetch-failure',
       defaultsTo: false,
@@ -70,6 +82,8 @@ Future<bool> run(List<String> arguments) async {
   final bool skipTemplate = parsedArguments['skip-template'] as bool;
   final bool verbose = parsedArguments['verbose'] as bool;
   final bool help = parsedArguments['help'] as bool;
+  final int numberShards = int.tryParse(parsedArguments['shards'] as String);
+  final int shardIndex = int.tryParse(parsedArguments['shard-index'] as String);
   final List<File> files = parsedArguments
     .rest
     .expand((String path) => Glob(path).listSync())
@@ -94,16 +108,38 @@ Future<bool> run(List<String> arguments) async {
   if (verbose)
     print('Starting run_tests.dart...');
 
+  if (files.length < shardIndex)
+    print('Warning: There are more shards than tests. Some shards will not run any tests.');
+
+  if (numberShards <= shardIndex) {
+    print('Error: There are more shard indexes than shards.');
+    return help;
+  }
+
+  // Best attempt at evenly splitting tests among the shards
+  final List<File> shardedFiles = <File>[];
+  for (int i = shardIndex; i < files.length; i += numberShards) {
+    shardedFiles.add(files[i]);
+  }
+
   int testCount = 0;
   int failures = 0;
 
   if (verbose) {
     final String s = files.length == 1 ? '' : 's';
-    print('${files.length} file$s specified.');
+    final String ss = shardedFiles.length == 1 ? '' : 's';
+    print('${files.length} file$s specified. ${shardedFiles.length} test$ss in shard #$shardIndex.');
     print('');
   }
 
-  for (final File file in files) {
+  if (verbose) {
+    print('Tests in this shard:');
+    for (final File file in shardedFiles)
+      print(file.path);
+  }
+  print('');
+
+  for (final File file in shardedFiles) {
     if (verbose)
       print('Processing ${file.path}...');
     TestFile instructions;