A simple Dart program for sharded test running

sharded_test_run.dart

import 'dart:async';
import 'dart:collection';
import 'dart:io';
import 'dart:math';

import 'package:args/args.dart';
import 'package:path/path.dart' as path;

ArgParser buildParser() => ArgParser()
  ..addOption('current-shard', abbr: 'c', mandatory: true)
  ..addOption('total-shards', abbr: 't', mandatory: true)
  ..addOption('seed', abbr: 's', mandatory: true)
  ..addOption('processes', abbr: 'p', defaultsTo: Platform.numberOfProcessors.toString());

Future<void> main(List<String> args) async {
  final parser = buildParser();

  final result = parser.parse(args);
  final currentShard = int.parse(result['current-shard']);
  final totalShards = int.parse(result['total-shards']);
  final seed = int.parse(result['seed']);
  final processes = int.parse(result['processes']);

  final packages = <Directory>[];
  final excludedPaths = ['pub-cache', '.dart_tool'];
  await for (final element in Directory.current.list(recursive: true)) {
    if (excludedPaths.any(element.path.contains)) continue;

    final pubspecExists = path.basename(element.path) == 'pubspec.yaml';
    final testFolderExists = Directory(path.join(element.parent.path, 'test')).existsSync();

    if (pubspecExists && testFolderExists) packages.add(element.parent);
  }

  stdout.writeln('Found ${packages.length} packages.');

  packages.shuffle(Random(seed));
  final shardSize = (packages.length / totalShards).ceil();
  final shards = [
    for (var i = 0; i < packages.length; i += shardSize)
      packages.sublist(i, min(i + shardSize, packages.length)),
  ];

  final thisShardPackages = shards.elementAtOrNull(currentShard - 1) ?? [];
  stdout.writeln('This shard runs ${thisShardPackages.length} packages.');

  final testsRunner = _TestsRunner(thisShardPackages, processes);
  final exitCodes = await testsRunner.run().then(
    (r) => r.entries.toList()..sort((a, b) => a.value.compareTo(b.value)),
  );

  for (final entry in exitCodes) {
    stdout.writeln('${path.basename(entry.key.path)}: ${entry.value != 0 ? '❌' : '✅'}');
  }

  final exitCode = exitCodes.any((entry) => entry.value != 0) ? 1 : 0;
  exit(exitCode);
}

class _TestsRunner {
  _TestsRunner(List<Directory> packages, this._processes) : _packageQueue = Queue.of(packages);

  final int _processes;
  final Queue<Directory> _packageQueue;
  final _exitCodes = <Directory, int>{};
  final _inProgress = <Future<int>>[];
  final _completer = Completer<void>();

  Future<Map<Directory, int>> run() async {
    unawaited(_processQueue());
    await _completer.future;

    return _exitCodes;
  }

  Future<void> _processQueue() async {
    while (_inProgress.length < _processes && _packageQueue.isNotEmpty) {
      final package = _packageQueue.removeFirst();
      final future = _runTestsInPackage(package);
      unawaited(_handleFuture(future, package));
    }

    if (_packageQueue.isEmpty && _inProgress.isEmpty) {
      _completer.complete();
    }
  }

  Future<void> _handleFuture(Future<int> future, Directory package) async {
    _inProgress.add(future);

    try {
      final exitCode = await future;
      _exitCodes[package] = exitCode;
      _inProgress.remove(future);
      unawaited(_processQueue());
    } catch (error) {
      stderr.writeln('Error running tests in ${path.basename(package.path)}: $error');
      _exitCodes[package] = -1;
      _inProgress.remove(future);
      unawaited(_processQueue());
    }
  }

  Future<int> _runTestsInPackage(Directory package) async {
    // dart format off
    final process = await Process.start('flutter', [
      'test',
      '--no-pub',
      '--coverage',
      '--file-reporter', 'json:test_report.json',
      '-x', 'golden',
    ], workingDirectory: package.path);
    // dart format on

    process.stdout.listen(stdout.add);
    process.stderr.listen(stderr.add);

    return await process.exitCode;
  }
}