[flutter_tools] allow flutter drive to take screenshots when sent a terminating signal (#114118)

Author: christopherfujino

packages/flutter_tools/lib/executable.dart

@@ -189,6 +189,7 @@ List<FlutterCommand> generateCommands({
     fileSystem: globals.fs,
     logger: globals.logger,
     platform: globals.platform,
+    signals: globals.signals,
   ),
   EmulatorsCommand(),
   FormatCommand(verboseHelp: verboseHelp),

packages/flutter_tools/lib/src/commands/drive.dart

@@ -12,8 +12,10 @@ import '../application_package.dart';
 import '../artifacts.dart';
 import '../base/common.dart';
 import '../base/file_system.dart';
+import '../base/io.dart';
 import '../base/logger.dart';
 import '../base/platform.dart';
+import '../base/signals.dart';
 import '../build_info.dart';
 import '../dart/package_map.dart';
 import '../device.dart';
@@ -48,9 +50,11 @@ class DriveCommand extends RunCommandBase {
   DriveCommand({
     bool verboseHelp = false,
     @visibleForTesting FlutterDriverFactory? flutterDriverFactory,
+    @visibleForTesting this.signalsToHandle = const <ProcessSignal>{ProcessSignal.sigint, ProcessSignal.sigterm},
     required FileSystem fileSystem,
     required Logger? logger,
     required Platform platform,
+    required this.signals,
   }) : _flutterDriverFactory = flutterDriverFactory,
        _fileSystem = fileSystem,
        _logger = logger,
@@ -149,6 +153,11 @@ class DriveCommand extends RunCommandBase {
           valueHelp: 'profile_memory.json');
   }
 
+  final Signals signals;
+
+  /// The [ProcessSignal]s that will lead to a screenshot being taken (if the option is provided).
+  final Set<ProcessSignal> signalsToHandle;
+
   // `pub` must always be run due to the test script running from source,
   // even if an application binary is used. Default to true unless the user explicitly
   // specified not to.
@@ -270,7 +279,7 @@ class DriveCommand extends RunCommandBase {
         );
       }
 
-      final int testResult = await driverService.startTest(
+      final Future<int> testResultFuture = driverService.startTest(
         testFile,
         stringsArg('test-arguments'),
         <String, String>{},
@@ -286,6 +295,13 @@ class DriveCommand extends RunCommandBase {
         androidEmulator: boolArgDeprecated('android-emulator'),
         profileMemory: stringArgDeprecated('profile-memory'),
       );
+      // If the test is sent a signal, take a screenshot before exiting
+      final Map<ProcessSignal, Object> screenshotTokens = _registerScreenshotCallbacks((ProcessSignal signal) async {
+        _logger!.printError('Caught $signal');
+        await _takeScreenshot(device);
+      });
+      final int testResult = await testResultFuture;
+      _unregisterScreenshotCallbacks(screenshotTokens);
       if (testResult != 0 && screenshot != null) {
         // Take a screenshot while the app is still running.
         await _takeScreenshot(device);
@@ -315,6 +331,21 @@ class DriveCommand extends RunCommandBase {
     return FlutterCommandResult.success();
   }
 
+  Map<ProcessSignal, Object> _registerScreenshotCallbacks(Function(ProcessSignal) callback) {
+    _logger!.printTrace('Registering signal handlers...');
+    final Map<ProcessSignal, Object> tokens = <ProcessSignal, Object>{};
+    for (final ProcessSignal signal in signalsToHandle) {
+      tokens[signal] = signals.addHandler(signal, callback);
+    }
+    return tokens;
+  }
+
+  void _unregisterScreenshotCallbacks(Map<ProcessSignal, Object> tokens) {
+    _logger!.printTrace('Unregistering signal handlers...');
+    for (final MapEntry<ProcessSignal, Object> entry in tokens.entries) {
+      signals.removeHandler(entry.key, entry.value);
+    }
+  }
   String? _getTestFile() {
     if (argResults!['driver'] != null) {
       return stringArgDeprecated('driver');

packages/flutter_tools/test/commands.shard/hermetic/drive_test.dart

@@ -2,12 +2,17 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+import 'dart:async';
+import 'dart:io' as io;
+
 import 'package:file/memory.dart';
 import 'package:flutter_tools/src/application_package.dart';
 import 'package:flutter_tools/src/base/common.dart';
 import 'package:flutter_tools/src/base/file_system.dart';
+import 'package:flutter_tools/src/base/io.dart';
 import 'package:flutter_tools/src/base/logger.dart';
 import 'package:flutter_tools/src/base/platform.dart';
+import 'package:flutter_tools/src/base/signals.dart';
 import 'package:flutter_tools/src/build_info.dart';
 import 'package:flutter_tools/src/cache.dart';
 import 'package:flutter_tools/src/commands/drive.dart';
@@ -27,12 +32,14 @@ void main() {
   late BufferLogger logger;
   late Platform platform;
   late FakeDeviceManager fakeDeviceManager;
+  late Signals signals;
 
   setUp(() {
     fileSystem = MemoryFileSystem.test();
     logger = BufferLogger.test();
     platform = FakePlatform();
     fakeDeviceManager = FakeDeviceManager();
+    signals = FakeSignals();
   });
 
   setUpAll(() {
@@ -44,7 +51,12 @@ void main() {
   });
 
   testUsingContext('warns if screenshot is not supported but continues test', () async {
-    final DriveCommand command = DriveCommand(fileSystem: fileSystem, logger: logger, platform: platform);
+    final DriveCommand command = DriveCommand(
+      fileSystem: fileSystem,
+      logger: logger,
+      platform: platform,
+      signals: signals,
+    );
     fileSystem.file('lib/main.dart').createSync(recursive: true);
     fileSystem.file('test_driver/main_test.dart').createSync(recursive: true);
     fileSystem.file('pubspec.yaml').createSync();
@@ -76,7 +88,12 @@ void main() {
   });
 
   testUsingContext('takes screenshot and rethrows on drive exception', () async {
-    final DriveCommand command = DriveCommand(fileSystem: fileSystem, logger: logger, platform: platform);
+    final DriveCommand command = DriveCommand(
+      fileSystem: fileSystem,
+      logger: logger,
+      platform: platform,
+      signals: signals,
+    );
     fileSystem.file('lib/main.dart').createSync(recursive: true);
     fileSystem.file('test_driver/main_test.dart').createSync(recursive: true);
     fileSystem.file('pubspec.yaml').createSync();
@@ -111,6 +128,7 @@ void main() {
       fileSystem: fileSystem,
       logger: logger,
       platform: platform,
+      signals: signals,
       flutterDriverFactory: FailingFakeFlutterDriverFactory(),
     );
 
@@ -149,7 +167,13 @@ void main() {
   });
 
   testUsingContext('drive --screenshot errors but does not fail if screenshot fails', () async {
-    final DriveCommand command = DriveCommand(fileSystem: fileSystem, logger: logger, platform: platform);
+    final DriveCommand command = DriveCommand(
+      fileSystem: fileSystem,
+      logger: logger,
+      platform: platform,
+      signals: signals,
+    );
+
     fileSystem.file('lib/main.dart').createSync(recursive: true);
     fileSystem.file('test_driver/main_test.dart').createSync(recursive: true);
     fileSystem.file('pubspec.yaml').createSync();
@@ -179,8 +203,71 @@ void main() {
     DeviceManager: () => fakeDeviceManager,
   });
 
+  testUsingContext('drive --screenshot takes screenshot if sent a registered signal', () async {
+    final FakeProcessSignal signal = FakeProcessSignal();
+    final ProcessSignal signalUnderTest = ProcessSignal(signal);
+    final DriveCommand command = DriveCommand(
+      fileSystem: fileSystem,
+      logger: logger,
+      platform: platform,
+      signals: Signals.test(),
+      flutterDriverFactory: NeverEndingFlutterDriverFactory(() {
+        signal.controller.add(signal);
+      }),
+      signalsToHandle: <ProcessSignal>{signalUnderTest},
+    );
+
+    fileSystem.file('lib/main.dart').createSync(recursive: true);
+    fileSystem.file('test_driver/main_test.dart').createSync(recursive: true);
+    fileSystem.file('pubspec.yaml').createSync();
+    fileSystem.directory('drive_screenshots').createSync();
+
+    final ScreenshotDevice screenshotDevice = ScreenshotDevice();
+    fakeDeviceManager.devices = <Device>[screenshotDevice];
+
+    expect(screenshotDevice.screenshots, isEmpty);
+
+    // This command will never complete. In reality, a real signal would have
+    // shut down the Dart process.
+    unawaited(
+      createTestCommandRunner(command).run(
+        <String>[
+          'drive',
+          '--no-pub',
+          '-d',
+          screenshotDevice.id,
+          '--use-existing-app',
+          'http://localhost:8181',
+          '--screenshot',
+          'drive_screenshots',
+        ],
+      ),
+    );
+
+    await screenshotDevice.firstScreenshot;
+    expect(
+      screenshotDevice.screenshots,
+      contains(isA<File>().having(
+        (File file) => file.path,
+        'path',
+        'drive_screenshots/drive_01.png',
+      )),
+    );
+  }, overrides: <Type, Generator>{
+    FileSystem: () => fileSystem,
+    ProcessManager: () => FakeProcessManager.any(),
+    Pub: () => FakePub(),
+    DeviceManager: () => fakeDeviceManager,
+  });
+
   testUsingContext('shouldRunPub is true unless user specifies --no-pub', () async {
-    final DriveCommand command = DriveCommand(fileSystem: fileSystem, logger: logger, platform: platform);
+    final DriveCommand command = DriveCommand(
+      fileSystem: fileSystem,
+      logger: logger,
+      platform: platform,
+      signals: signals,
+    );
+
     fileSystem.file('lib/main.dart').createSync(recursive: true);
     fileSystem.file('test_driver/main_test.dart').createSync(recursive: true);
     fileSystem.file('pubspec.yaml').createSync();
@@ -207,7 +294,13 @@ void main() {
   });
 
   testUsingContext('flags propagate to debugging options', () async {
-    final DriveCommand command = DriveCommand(fileSystem: fileSystem, logger: logger, platform: platform);
+    final DriveCommand command = DriveCommand(
+      fileSystem: fileSystem,
+      logger: logger,
+      platform: platform,
+      signals: signals,
+    );
+
     fileSystem.file('lib/main.dart').createSync(recursive: true);
     fileSystem.file('test_driver/main_test.dart').createSync(recursive: true);
     fileSystem.file('pubspec.yaml').createSync();
@@ -270,6 +363,13 @@ class ThrowingScreenshotDevice extends ScreenshotDevice {
 // Until we fix that, we have to also ignore related lints here.
 // ignore: avoid_implementing_value_types
 class ScreenshotDevice extends Fake implements Device {
+  final List<File> screenshots = <File>[];
+
+  final Completer<void> _firstScreenshotCompleter = Completer<void>();
+
+  /// A Future that completes when [takeScreenshot] is called the first time.
+  Future<void> get firstScreenshot => _firstScreenshotCompleter.future;
+
   @override
   final String name = 'FakeDevice';
 
@@ -299,7 +399,12 @@ class ScreenshotDevice extends Fake implements Device {
     }) async => LaunchResult.succeeded();
 
   @override
-  Future<void> takeScreenshot(File outputFile) async {}
+  Future<void> takeScreenshot(File outputFile) async {
+    if (!_firstScreenshotCompleter.isCompleted) {
+      _firstScreenshotCompleter.complete();
+    }
+    screenshots.add(outputFile);
+  }
 }
 
 class FakePub extends Fake implements Pub {
@@ -331,6 +436,48 @@ class FakeDeviceManager extends Fake implements DeviceManager {
   Future<List<Device>> findTargetDevices(FlutterProject? flutterProject, {Duration? timeout, bool promptUserToChooseDevice = true}) async => devices;
 }
 
+/// A [FlutterDriverFactory] that creates a [NeverEndingDriverService].
+class NeverEndingFlutterDriverFactory extends Fake implements FlutterDriverFactory {
+  NeverEndingFlutterDriverFactory(this.callback);
+
+  final void Function() callback;
+
+  @override
+  DriverService createDriverService(bool web) => NeverEndingDriverService(callback);
+}
+
+/// A [DriverService] that will return a Future from [startTest] that will never complete.
+///
+/// This is to similate when the test will take a long time, but a signal is
+/// expected to interrupt the process.
+class NeverEndingDriverService extends Fake implements DriverService {
+  NeverEndingDriverService(this.callback);
+
+  final void Function() callback;
+  @override
+  Future<void> reuseApplication(Uri vmServiceUri, Device device, DebuggingOptions debuggingOptions, bool ipv6) async { }
+
+  @override
+  Future<int> startTest(
+    String testFile,
+    List<String> arguments,
+    Map<String, String> environment,
+    PackageConfig packageConfig, {
+      bool? headless,
+      String? chromeBinary,
+      String? browserName,
+      bool? androidEmulator,
+      int? driverPort,
+      List<String>? webBrowserFlags,
+      List<String>? browserDimension,
+      String? profileMemory,
+    }) async {
+      callback();
+      // return a Future that will never complete.
+      return Completer<int>().future;
+  }
+}
+
 class FailingFakeFlutterDriverFactory extends Fake implements FlutterDriverFactory {
   @override
   DriverService createDriverService(bool web) => FailingFakeDriverService();
@@ -356,3 +503,10 @@ class FailingFakeDriverService extends Fake implements DriverService {
       String? profileMemory,
     }) async => 1;
 }
+
+class FakeProcessSignal extends Fake implements io.ProcessSignal {
+  final StreamController<io.ProcessSignal> controller = StreamController<io.ProcessSignal>();
+
+  @override
+  Stream<io.ProcessSignal> watch() => controller.stream;
+}