feat(test): support sharding tests across CI runners

bricks/test_optimizer/brick.yaml

@@ -12,3 +12,11 @@ vars:
     default: "."
     description: The path to the package root.
     prompt: Please enter the path to the package root.
+  shard-index:
+    type: number
+    description: The 1-based index of the shard to generate tests for.
+    prompt: Please enter the shard index.
+  total-shards:
+    type: number
+    description: The total number of shards the test suite is split into.
+    prompt: Please enter the total number of shards.

bricks/test_optimizer/hooks/lib/pre_gen.dart

@@ -33,35 +33,96 @@ Future<void> run(HookContext context) async {
   final flutterSdkRegExp = RegExp(r'sdk:\s*flutter$', multiLine: true);
   final isFlutter = flutterSdkRegExp.hasMatch(pubspecContents);
 
-  final identifierGenerator = DartIdentifierGenerator();
-  final testIdentifierTable = <Map<String, String>>[];
+  final shardIndex = context.vars['shard-index'] as int?;
+  final totalShards = context.vars['total-shards'] as int?;
+
+  // The CLI validates these before it gets here, but `mason make` prompts for
+  // them directly, so guard the round-robin below against values that would
+  // never terminate or index out of range.
+  if (shardIndex != null &&
+      totalShards != null &&
+      (totalShards < 1 || shardIndex < 1 || shardIndex > totalShards)) {
+    context.logger.err(
+      'shard-index must be between 1 and total-shards, but got '
+      'shard-index $shardIndex and total-shards $totalShards',
+    );
+    exitFn(1);
+  }
+
   final tests = testDir
       .listSync(recursive: true)
       .where((entity) => entity.isTest);
 
-  final notOptimizedTests = await getNotOptimizedTests(tests, testDir.path);
-
-  for (final entity in tests) {
-    final relativePath = path
-        .relative(entity.path, from: testDir.path)
-        .replaceAll(r'\', '/');
-    testIdentifierTable.add({
-      'path': relativePath,
-      'identifier': identifierGenerator.next(),
-    });
-  }
-
-  final optimizedTestsIdentifierTable = testIdentifierTable
-      .where((e) => !notOptimizedTests.contains(e['path']))
+  final notOptimizedTests = (await getNotOptimizedTests(
+    tests,
+    testDir.path,
+  )).toSet();
+
+  // Sorting guarantees a deterministic order across machines, which is what
+  // makes sharding reproducible: `Directory.listSync` order is filesystem
+  // dependent, so without this two runners could disagree on the partition
+  // and either skip or duplicate tests.
+  final testPaths =
+      tests
+          .map(
+            (entity) => path
+                .relative(entity.path, from: testDir.path)
+                .replaceAll(r'\', '/'),
+          )
+          .toList()
+        ..sort();
+
+  // Non optimized tests run as standalone files alongside the optimizer
+  // entrypoint, so they are sharded too, and in the same deal as the
+  // optimized ones: dealing out one list keeps every shard within one file
+  // of the others, whereas dealing out the two lists separately would hand
+  // the first shards a file from each.
+  final shardPaths = _shardOf(
+    testPaths,
+    shardIndex: shardIndex,
+    totalShards: totalShards,
+  );
+  final optimizedTestPaths = shardPaths
+      .where((p) => !notOptimizedTests.contains(p))
+      .toList();
+  final shardedNotOptimizedTests = shardPaths
+      .where(notOptimizedTests.contains)
       .toList();
 
+  final identifierGenerator = DartIdentifierGenerator();
+  final optimizedTestsIdentifierTable = [
+    for (final relativePath in optimizedTestPaths)
+      {'path': relativePath, 'identifier': identifierGenerator.next()},
+  ];
+
   context.vars = {
     'tests': optimizedTestsIdentifierTable,
     'isFlutter': isFlutter,
-    'notOptimizedTests': notOptimizedTests,
+    'notOptimizedTests': shardedNotOptimizedTests,
   };
 }
 
+/// Returns the subset of [paths] that belongs to the shard [shardIndex] out of
+/// [totalShards].
+///
+/// Returns [paths] unchanged when sharding is not enabled (either value is
+/// `null`).
+///
+/// Files are dealt out round-robin (index modulo [totalShards]) over the
+/// already sorted [paths], which keeps shards balanced in file count and makes
+/// the partition stable for a given test suite.
+List<String> _shardOf(
+  List<String> paths, {
+  required int? shardIndex,
+  required int? totalShards,
+}) {
+  if (shardIndex == null || totalShards == null) return paths;
+
+  return [
+    for (var i = shardIndex - 1; i < paths.length; i += totalShards) paths[i],
+  ];
+}
+
 extension on FileSystemEntity {
   bool get isTest {
     return this is File && path.basename(this.path).endsWith('_test.dart');
@@ -85,9 +146,10 @@ Future<List<String>> getNotOptimizedTests(
     }
   }
 
-  /// Format to relative path
+  /// Format to relative path, normalizing separators so the paths compare
+  /// equal to the ones built in [run] on Windows too.
   final relativePaths = testWithVeryGoodTest
-      .map((e) => path.relative(e, from: testDir))
+      .map((e) => path.relative(e, from: testDir).replaceAll(r'\', '/'))
       .toList();
 
   return relativePaths;

bricks/test_optimizer/hooks/test/pre_gen_test.dart

@@ -1,4 +1,5 @@
 import 'dart:io';
+import 'dart:math';
 
 import 'package:hooks/pre_gen.dart' as pre_gen;
 import 'package:mason/mason.dart';
@@ -188,6 +189,33 @@ dependencies:
         expect(context.vars['isFlutter'], isNull);
       });
 
+      test('when the shard values are out of range', () async {
+        File(path.join(tempDirectory.path, 'pubspec.yaml')).createSync();
+        Directory(path.join(tempDirectory.path, 'test')).createSync();
+
+        context.vars['package-root'] = tempDirectory.absolute.path;
+        context.vars['shard-index'] = 1;
+        context.vars['total-shards'] = 0;
+
+        await expectLater(
+          () => pre_gen.run(context),
+          throwsA(
+            isA<ProcessException>().having(
+              (ex) => ex.arguments.first,
+              'error code',
+              equals('1'),
+            ),
+          ),
+        );
+
+        verify(
+          () => context.logger.err(
+            'shard-index must be between 1 and total-shards, but got '
+            'shard-index 1 and total-shards 0',
+          ),
+        ).called(1);
+      });
+
       test('when target dir does not contain a pubspec.yaml', () async {
         final testDir = Directory(path.join(tempDirectory.path, 'test'))
           ..createSync();
@@ -278,5 +306,142 @@ dependencies:
         },
       );
     });
+    group('Sharding', () {
+      /// Creates a package with [count] optimizable test files, plus any
+      /// [notOptimized] files carrying the skip optimization tag.
+      Directory createPackage(int count, {int notOptimized = 0}) {
+        File(path.join(tempDirectory.path, 'pubspec.yaml')).createSync();
+        final testDir = Directory(path.join(tempDirectory.path, 'test'))
+          ..createSync();
+        for (var i = 0; i < count; i++) {
+          File(path.join(testDir.path, 'test${i}_test.dart')).createSync();
+        }
+        for (var i = 0; i < notOptimized; i++) {
+          File(path.join(testDir.path, 'skip${i}_test.dart'))
+              .writeAsStringSync(notOptimizedTestContent);
+        }
+        return testDir;
+      }
+
+      List<String> pathsOf(HookContext context) {
+        final tests = context.vars['tests'] as List<Map<String, String>>;
+        return tests.map((e) => e['path']!).toList();
+      }
+
+      Future<List<String>> runShard(int index, int total) async {
+        final context = _FakeContext()
+          ..vars['package-root'] = tempDirectory.absolute.path
+          ..vars['shard-index'] = index
+          ..vars['total-shards'] = total;
+        await pre_gen.run(context);
+        return [
+          ...pathsOf(context),
+          ...(context.vars['notOptimizedTests']! as List).cast<String>(),
+        ];
+      }
+
+      test('runs every test exactly once across all shards', () async {
+        createPackage(7, notOptimized: 2);
+
+        final shards = [for (var i = 1; i <= 3; i++) await runShard(i, 3)];
+        final union = shards.expand((shard) => shard).toList();
+
+        expect(
+          union..sort(),
+          [
+            for (var i = 0; i < 7; i++) 'test${i}_test.dart',
+            for (var i = 0; i < 2; i++) 'skip${i}_test.dart',
+          ]..sort(),
+          reason: 'Shards must be a complete and disjoint partition',
+        );
+      });
+
+      test('shards non optimized tests as well', () async {
+        createPackage(0, notOptimized: 4);
+
+        final first = await runShard(1, 2);
+        final second = await runShard(2, 2);
+
+        expect(first, ['skip0_test.dart', 'skip2_test.dart']);
+        expect(second, ['skip1_test.dart', 'skip3_test.dart']);
+      });
+
+      test('deals optimized and non optimized tests out together', () async {
+        createPackage(2, notOptimized: 3);
+
+        final sizes = [
+          for (var i = 1; i <= 6; i++) (await runShard(i, 6)).length,
+        ];
+
+        expect(
+          sizes,
+          [1, 1, 1, 1, 1, 0],
+          reason:
+              'Sharding the two lists separately would give the first '
+              'shards a file from each while later shards stay empty',
+        );
+      });
+
+      test('is deterministic across runs', () async {
+        createPackage(9);
+
+        expect(await runShard(2, 4), await runShard(2, 4));
+      });
+
+      test('balances shards within one file of each other', () async {
+        createPackage(10);
+
+        final sizes = [
+          for (var i = 1; i <= 4; i++) (await runShard(i, 4)).length,
+        ];
+
+        expect(sizes.reduce(max) - sizes.reduce(min), lessThanOrEqualTo(1));
+      });
+
+      test(
+        'yields an empty shard when there are more shards than tests',
+        () async {
+          createPackage(2);
+
+          expect(await runShard(3, 3), isEmpty);
+        },
+      );
+
+      test(
+        'excludes nested non optimized tests from the optimized set',
+        () async {
+          File(path.join(tempDirectory.path, 'pubspec.yaml')).createSync();
+          final testDir = Directory(path.join(tempDirectory.path, 'test'))
+            ..createSync();
+          final nested = Directory(path.join(testDir.path, 'sub'))
+            ..createSync();
+          File(path.join(nested.path, 'skip_test.dart'))
+              .writeAsStringSync(notOptimizedTestContent);
+
+          final context = _FakeContext()
+            ..vars['package-root'] = tempDirectory.absolute.path;
+          await pre_gen.run(context);
+
+          expect(
+            pathsOf(context),
+            isEmpty,
+            reason:
+                'A tagged test in a subdirectory must not be optimized, '
+                'otherwise it runs both inlined and standalone',
+          );
+          expect(context.vars['notOptimizedTests'], ['sub/skip_test.dart']);
+        },
+      );
+
+      test('includes every test when sharding is not requested', () async {
+        createPackage(3);
+
+        final context = _FakeContext()
+          ..vars['package-root'] = tempDirectory.absolute.path;
+        await pre_gen.run(context);
+
+        expect(pathsOf(context), hasLength(3));
+      });
+    });
   });
 }

lib/src/cli/dart_cli.dart

@@ -136,6 +136,8 @@ class Dart {
     void Function(String)? stderr,
     GeneratorBuilder buildGenerator = MasonGenerator.fromBundle,
     List<String>? reportOn,
+    int? shardIndex,
+    int? totalShards,
   }) {
     return TestCLIRunner.test(
       logger: logger,
@@ -157,6 +159,8 @@ class Dart {
       stderr: stderr,
       reportOn: reportOn,
       buildGenerator: buildGenerator,
+      shardIndex: shardIndex,
+      totalShards: totalShards,
     );
   }
 }

lib/src/cli/flutter_cli.dart

@@ -215,6 +215,8 @@ class Flutter {
     void Function(String)? stderr,
     GeneratorBuilder buildGenerator = MasonGenerator.fromBundle,
     List<String>? reportOn,
+    int? shardIndex,
+    int? totalShards,
   }) {
     return TestCLIRunner.test(
       logger: logger,
@@ -235,6 +237,8 @@ class Flutter {
       stderr: stderr,
       buildGenerator: buildGenerator,
       reportOn: reportOn,
+      shardIndex: shardIndex,
+      totalShards: totalShards,
     );
   }
 }