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

@@ -40,6 +40,22 @@ Future<void> run(HookContext context) async {
   final flutterSdkRegExp = RegExp(r'sdk:\s*flutter$', multiLine: true);
   final isFlutter = flutterSdkRegExp.hasMatch(pubspecContents);
 
+  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 identifierGenerator = DartIdentifierGenerator();
   final optimizedTests = <Map<String, String>>[];
   final notOptimizedTests = <String>[];
@@ -48,13 +64,23 @@ Future<void> run(HookContext context) async {
       .listSync(recursive: true)
       .where((entity) => entity.isTest)
       .cast<File>();
-  final parsedTests = await Future.wait(tests.map(_parse));
-
-  for (final (file, content, metadata) in parsedTests) {
-    final relativePath = path
-        .relative(file.path, from: testDir.path)
-        .replaceAll(r'\', '/');
-
+  final parsedTests = await Future.wait(
+    tests.map((file) => _parse(file, testDir: testDir.path)),
+  );
+
+  // 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. Tests kept out of the bundle are
+  // dealt out in the same round as the optimized ones, which keeps every
+  // shard within one file of the others.
+  final shard = _shardOf(
+    parsedTests..sort((a, b) => a.relativePath.compareTo(b.relativePath)),
+    shardIndex: shardIndex,
+    totalShards: totalShards,
+  );
+
+  for (final (:relativePath, :content, :metadata) in shard) {
     if (metadata.skipsOptimization) {
       notOptimizedTests.add(relativePath);
       continue;
@@ -89,11 +115,39 @@ Future<void> run(HookContext context) async {
   };
 }
 
-typedef _ParsedTest = (File file, String content, TestMetadata metadata);
+typedef _ParsedTest = ({
+  String relativePath,
+  String content,
+  TestMetadata metadata,
+});
 
-Future<_ParsedTest> _parse(File file) async {
+/// Reads and parses [file], keyed by its POSIX path relative to [testDir].
+Future<_ParsedTest> _parse(File file, {required String testDir}) async {
   final content = await file.readAsString();
-  return (file, content, parseTestMetadata(content, path: file.path));
+  return (
+    relativePath: path.relative(file.path, from: testDir).replaceAll(r'\', '/'),
+    content: content,
+    metadata: parseTestMetadata(content, path: file.path),
+  );
+}
+
+/// Returns the subset of [items] that belongs to the shard [shardIndex] out of
+/// [totalShards], or [items] unchanged when sharding is not enabled (either
+/// value is `null`).
+///
+/// Items are dealt out round-robin (index modulo [totalShards]) over the
+/// already sorted [items], which keeps shards balanced in file count and makes
+/// the partition stable for a given test suite.
+List<T> _shardOf<T>(
+  List<T> items, {
+  required int? shardIndex,
+  required int? totalShards,
+}) {
+  if (shardIndex == null || totalShards == null) return items;
+
+  return [
+    for (var i = shardIndex - 1; i < items.length; i += totalShards) items[i],
+  ];
 }
 
 extension on FileSystemEntity {

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:hooks/test_metadata.dart';
@@ -305,6 +306,33 @@ void main() {}
         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();
@@ -412,5 +440,142 @@ void main() {}
         expect(metadata.skipsOptimization, isFalse);
       });
     });
+    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

@@ -137,6 +137,8 @@ class Dart {
     void Function(String)? stderr,
     GeneratorBuilder buildGenerator = MasonGenerator.fromBundle,
     List<String>? reportOn,
+    int? shardIndex,
+    int? totalShards,
   }) {
     return TestCLIRunner.test(
       logger: logger,
@@ -150,6 +152,8 @@ class Dart {
         enabled: optimizePerformance,
         exclude: excludeOptimization,
         buildGenerator: buildGenerator,
+        shardIndex: shardIndex,
+        totalShards: totalShards,
       ),
       ignore: ignore,
       minCoverage: minCoverage,

lib/src/cli/flutter_cli.dart

@@ -214,6 +214,8 @@ class Flutter {
     void Function(String)? stderr,
     GeneratorBuilder buildGenerator = MasonGenerator.fromBundle,
     List<String>? reportOn,
+    int? shardIndex,
+    int? totalShards,
   }) {
     return TestCLIRunner.test(
       logger: logger,
@@ -225,6 +227,8 @@ class Flutter {
         enabled: optimizePerformance,
         exclude: excludeOptimization,
         buildGenerator: buildGenerator,
+        shardIndex: shardIndex,
+        totalShards: totalShards,
       ),
       ignore: ignore,
       minCoverage: minCoverage,