[flutter_driver] use mostly public screenshot API. (#157888)

Instead of completely private. This has been broken for Impeller for years, which shows how much this method is getting used.

Fixes https://github.com/flutter/flutter/issues/130461

.ci.yaml

@@ -1120,6 +1120,23 @@ targets:
       - bin/**
       - .ci.yaml
 
+  - name: Linux linux_desktop_impeller
+    recipe: devicelab/devicelab_drone
+    timeout: 60
+    presubmit: false
+    bringup: true
+    properties:
+      xvfb: "1"
+      dependencies: >-
+        [
+          {"dependency": "clang", "version": "git_revision:5d5aba78dbbee75508f01bcaa69aedb2ab79065a"},
+          {"dependency": "cmake", "version": "build_id:8787856497187628321"},
+          {"dependency": "ninja", "version": "version:1.9.0"}
+        ]
+      tags: >
+        ["devicelab", "hostonly", "linux"]
+      task_name: linux_desktop_impeller
+
   - name: Linux run_release_test_linux
     recipe: devicelab/devicelab_drone
     timeout: 60
@@ -5363,6 +5380,20 @@ targets:
       - bin/**
       - .ci.yaml
 
+  - name: Mac_arm64 mac_desktop_impeller
+    recipe: devicelab/devicelab_drone
+    presubmit: false
+    bringup: true
+    timeout: 60
+    properties:
+      dependencies: >-
+        [
+          {"dependency": "ruby", "version": "ruby_3.1-pod_1.13"}
+        ]
+      tags: >
+        ["devicelab", "hostonly", "mac", "arm64"]
+      task_name: run_release_test_macos
+
   - name: Windows build_tests_1_8
     recipe: flutter/flutter_drone
     timeout: 60
@@ -6295,6 +6326,20 @@ targets:
         ]
       task_name: hello_world_win_desktop__compile
 
+  - name: Windows windows_desktop_impeller
+    recipe: devicelab/devicelab_drone
+    bringup: true
+    presubmit: false
+    timeout: 60
+    properties:
+      tags: >
+        ["devicelab", "hostonly", "windows"]
+      dependencies: >-
+        [
+          {"dependency": "vs_build", "version": "version:vs2019"}
+        ]
+      task_name: windows_desktop_impeller
+
   - name: Windows_arm64 hello_world_win_desktop__compile
     recipe: devicelab/devicelab_drone
     presubmit: false

TESTOWNERS

@@ -294,6 +294,9 @@
 /dev/devicelab/bin/tasks/web_benchmarks_skwasm.dart @eyebrowsoffire @flutter/web
 /dev/devicelab/bin/tasks/windows_home_scroll_perf__timeline_summary.dart @jonahwilliams @flutter/engine
 /dev/devicelab/bin/tasks/windows_startup_test.dart @loic-sharma @flutter/desktop
+/dev/devicelab/bin/tasks/windows_desktop_impeller.dart @jonahwilliams @flutter/engine
+/dev/devicelab/bin/tasks/mac_desktop_impeller.dart @jonahwilliams @flutter/engine
+/dev/devicelab/bin/tasks/linux_desktop_impeller.dart @jonahwilliams @flutter/engine
 
 ## Host only framework tests
 # Linux docs_deploy_beta

dev/devicelab/bin/tasks/linux_desktop_impeller.dart

@@ -0,0 +1,12 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter_devicelab/framework/devices.dart';
+import 'package:flutter_devicelab/framework/framework.dart';
+import 'package:flutter_devicelab/tasks/integration_tests.dart';
+
+Future<void> main() async {
+  deviceOperatingSystem = DeviceOperatingSystem.linux;
+  await task(createSolidColorTest(enableImpeller: true));
+}

dev/devicelab/bin/tasks/mac_desktop_impeller.dart

@@ -0,0 +1,12 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter_devicelab/framework/devices.dart';
+import 'package:flutter_devicelab/framework/framework.dart';
+import 'package:flutter_devicelab/tasks/integration_tests.dart';
+
+Future<void> main() async {
+  deviceOperatingSystem = DeviceOperatingSystem.macos;
+  await task(createSolidColorTest(enableImpeller: true));
+}

dev/devicelab/bin/tasks/windows_desktop_impeller.dart

@@ -0,0 +1,12 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter_devicelab/framework/devices.dart';
+import 'package:flutter_devicelab/framework/framework.dart';
+import 'package:flutter_devicelab/tasks/integration_tests.dart';
+
+Future<void> main() async {
+  deviceOperatingSystem = DeviceOperatingSystem.windows;
+  await task(createSolidColorTest(enableImpeller: true));
+}

dev/devicelab/lib/tasks/integration_tests.dart

@@ -131,6 +131,17 @@ TaskFunction createEndToEndKeyboardTextfieldTest() {
   ).call;
 }
 
+TaskFunction createSolidColorTest({required bool enableImpeller}) {
+  return DriverTest(
+    '${flutterDirectory.path}/dev/integration_tests/ui',
+    'lib/solid_color.dart',
+    extraOptions: <String>[
+      if (enableImpeller)
+        '--enable-impeller'
+    ]
+  ).call;
+}
+
 TaskFunction dartDefinesTask() {
   return DriverTest(
     '${flutterDirectory.path}/dev/integration_tests/ui',

dev/integration_tests/ui/lib/solid_color.dart

@@ -0,0 +1,14 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter/material.dart';
+import 'package:flutter_driver/driver_extension.dart';
+
+void main() {
+  enableFlutterDriverExtension();
+
+  runApp(Container(
+    color: const Color(0xFFFF0000),
+  ));
+}

dev/integration_tests/ui/test_driver/solid_color_test.dart

@@ -0,0 +1,28 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'dart:typed_data';
+
+import 'package:flutter_driver/flutter_driver.dart';
+import 'package:test/test.dart' hide TypeMatcher, isInstanceOf;
+
+void main() {
+  late FlutterDriver driver;
+
+  setUpAll(() async {
+    driver = await FlutterDriver.connect();
+    await driver.waitUntilFirstFrameRasterized();
+  });
+
+  test('Can render solid red', () async {
+    // RGBA Encoded Bytes.
+    final Uint8List data = (await driver.screenshot(format: ScreenshotFormat.rawStraightRgba)) as Uint8List;
+
+    expect(data[0] << 24 | data[1] << 16 | data[2] << 8 | data[3], 0xFF0000FF);
+  }, timeout: Timeout.none);
+
+  tearDownAll(() async {
+    await driver.close();
+  });
+}

packages/flutter_driver/lib/src/common/deserialization_factory.dart

@@ -13,6 +13,7 @@ import 'layer_tree.dart';
 import 'message.dart';
 import 'render_tree.dart';
 import 'request_data.dart';
+import 'screenshot.dart';
 import 'semantics.dart';
 import 'text.dart';
 import 'text_input_action.dart';
@@ -64,6 +65,7 @@ mixin DeserializeCommandFactory {
       'get_semantics_id'              => GetSemanticsId.deserialize(params, finderFactory),
       'get_offset'                    => GetOffset.deserialize(params, finderFactory),
       'get_diagnostics_tree'          => GetDiagnosticsTree.deserialize(params, finderFactory),
+      'screenshot'                    => ScreenshotCommand.deserialize(params),
       final String? kind => throw DriverError('Unsupported command kind $kind'),
     };
   }

packages/flutter_driver/lib/src/common/handler_factory.dart

@@ -6,6 +6,7 @@
 library;
 
 import 'dart:async';
+import 'dart:ui' as ui;
 
 import 'package:flutter/cupertino.dart';
 import 'package:flutter/foundation.dart';
@@ -27,6 +28,7 @@ import 'layer_tree.dart';
 import 'message.dart';
 import 'render_tree.dart';
 import 'request_data.dart';
+import 'screenshot.dart';
 import 'semantics.dart';
 import 'text.dart';
 import 'text_input_action.dart' show SendTextInputAction;
@@ -164,6 +166,7 @@ mixin CommandHandlerFactory {
       'get_semantics_id'              => _getSemanticsId(command, finderFactory),
       'get_offset'                    => _getOffset(command, finderFactory),
       'get_diagnostics_tree'          => _getDiagnosticsTree(command, finderFactory),
+      'screenshot'                    => _takeScreenshot(command),
       final String kind => throw DriverError('Unsupported command kind $kind'),
     };
 
@@ -352,6 +355,18 @@ mixin CommandHandlerFactory {
     )));
   }
 
+  Future<ScreenshotResult> _takeScreenshot(Command command) async {
+    final ScreenshotCommand screenshotCommand = command as ScreenshotCommand;
+    final RenderView renderView = RendererBinding.instance.renderViews.first;
+    // ignore: invalid_use_of_protected_member
+    final ContainerLayer? layer = renderView.layer;
+    final OffsetLayer offsetLayer = layer! as OffsetLayer;
+    final ui.Image image = await offsetLayer.toImage(renderView.paintBounds);
+    final ui.ImageByteFormat format = ui.ImageByteFormat.values[screenshotCommand.format.index];
+    final ByteData buffer = (await image.toByteData(format: format))!;
+    return ScreenshotResult(buffer.buffer.asUint8List());
+  }
+
   Future<Result> _scroll(Command command, WidgetController prober, CreateFinderFactory finderFactory) async {
     final Scroll scrollCommand = command as Scroll;
     final Finder target = await waitForElement(finderFactory.createFinder(scrollCommand.finder));

packages/flutter_driver/lib/src/common/screenshot.dart

@@ -0,0 +1,72 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'dart:convert';
+import 'dart:typed_data';
+
+import 'message.dart';
+
+/// Format of data returned by screenshot driver command.
+enum ScreenshotFormat {
+  /// Raw RGBA format.
+  ///
+  /// Unencoded bytes, in RGBA row-primary form with premultiplied alpha, 8 bits per channel.
+  rawRgba,
+
+  /// Raw straight RGBA format.
+  ///
+  /// Unencoded bytes, in RGBA row-primary form with straight alpha, 8 bits per channel.
+  rawStraightRgba,
+
+  /// Raw unmodified format.
+  rawUnmodified,
+
+  /// Raw extended range RGBA format.
+  ///
+  /// Unencoded bytes, in RGBA row-primary form with straight alpha, 32 bit
+  /// float (IEEE 754 binary32) per channel.
+  rawExtendedRgba128,
+
+  /// PNG format.
+  png,
+}
+
+/// A Flutter Driver command that takes a screenshot.
+class ScreenshotCommand extends Command {
+  /// Constructs this command given a [finder].
+  ScreenshotCommand({super.timeout, this.format = ScreenshotFormat.png});
+
+  /// Deserializes this command from the value generated by [serialize].
+  ScreenshotCommand.deserialize(super.json)
+      : format = ScreenshotFormat.values[int.tryParse(json['format']!) ?? 4], super.deserialize();
+
+  /// Whether the resulting data is PNG compressed.
+  final ScreenshotFormat format;
+
+  /// Serializes this command to parameter name/value pairs.
+  @override
+  Map<String, String> serialize() {
+    return super.serialize()..addAll(<String, String>{
+      'format': format.index.toString()
+    });
+  }
+
+  @override
+  String get kind => 'screenshot';
+}
+
+/// base64 encode a PNG
+class ScreenshotResult extends Result {
+  /// Consructs a screenshot result with PNG or raw RGBA byte data.
+  ScreenshotResult(this._data);
+
+  final Uint8List _data;
+
+  @override
+  Map<String, Object?> toJson() {
+    return <String, Object?>{
+      'data': base64.encode(_data)
+    };
+  }
+}

packages/flutter_driver/lib/src/driver/driver.dart

@@ -9,6 +9,7 @@
 /// @docImport 'package:flutter_test/flutter_test.dart';
 library;
 
+import 'dart:convert';
 import 'dart:io';
 
 import 'package:meta/meta.dart';
@@ -26,6 +27,7 @@ import '../common/layer_tree.dart';
 import '../common/message.dart';
 import '../common/render_tree.dart';
 import '../common/request_data.dart';
+import '../common/screenshot.dart';
 import '../common/semantics.dart';
 import '../common/text.dart';
 import '../common/text_input_action.dart';
@@ -34,6 +36,7 @@ import 'timeline.dart';
 import 'vmservice_driver.dart';
 import 'web_driver.dart';
 
+export '../common/screenshot.dart' show ScreenshotFormat;
 export 'vmservice_driver.dart';
 export 'web_driver.dart';
 
@@ -630,8 +633,10 @@ abstract class FlutterDriver {
   /// In practice, sometimes the device gets really busy for a while and even
   /// two seconds isn't enough, which means that this is still racy and a source
   /// of flakes.
-  Future<List<int>> screenshot() async {
-    throw UnimplementedError();
+  Future<List<int>> screenshot({ScreenshotFormat format = ScreenshotFormat.png}) async {
+    await Future<void>.delayed(const Duration(seconds: 2));
+    final Map<String, Object?> jsonResponse = await sendCommand(ScreenshotCommand(format: format));
+    return base64.decode(jsonResponse['data']! as String);
   }
 
   /// Returns the Flags set in the Dart VM as JSON.

packages/flutter_driver/lib/src/driver/vmservice_driver.dart

@@ -6,7 +6,6 @@
 library;
 
 import 'dart:async';
-import 'dart:convert';
 import 'dart:io';
 
 import 'package:file/file.dart' as f;
@@ -359,14 +358,6 @@ class VMServiceFlutterDriver extends FlutterDriver {
     }
   }
 
-  @override
-  Future<List<int>> screenshot() async {
-    await Future<void>.delayed(const Duration(seconds: 2));
-
-    final vms.Response result = await _serviceClient.callMethod('_flutter.screenshot');
-    return base64.decode(result.json!['screenshot'] as String);
-  }
-
   @override
   Future<List<Map<String, dynamic>>> getVmFlags() async {
     final vms.FlagList result = await _serviceClient.getFlagList();

packages/flutter_driver/lib/src/driver/web_driver.dart

@@ -181,7 +181,10 @@ class WebFlutterDriver extends FlutterDriver {
   }
 
   @override
-  Future<List<int>> screenshot() async {
+  Future<List<int>> screenshot({ScreenshotFormat format = ScreenshotFormat.png}) async {
+    if (format != ScreenshotFormat.png) {
+      throw ArgumentError.value(format, 'format', 'Web Driver only supports PNG screenshot format');
+    }
     await Future<void>.delayed(const Duration(seconds: 2));
 
     return _connection.screenshot();

packages/flutter_driver/test_driver/screenshot.dart

@@ -0,0 +1,11 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'package:flutter/material.dart';
+import 'package:flutter_driver/driver_extension.dart';
+
+void main() {
+  enableFlutterDriverExtension();
+  runApp(Container(color: Colors.red));
+}

packages/flutter_driver/test_driver/screenshot_test.dart

@@ -0,0 +1,30 @@
+// Copyright 2014 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+import 'dart:typed_data';
+
+import 'package:flutter_driver/flutter_driver.dart';
+
+import '../test/common.dart';
+
+void main() {
+  late FlutterDriver driver;
+
+  setUpAll(() async {
+    driver = await FlutterDriver.connect();
+    await driver.waitUntilFirstFrameRasterized();
+  });
+
+  test('it takes a screenshot', () async {
+    // PNG Encoded Bytes.
+    final Uint8List bytes = (await driver.screenshot()) as Uint8List;
+
+    // Check PNG header.
+    expect(bytes.sublist(0, 8), <int>[137, 80, 78, 71, 13, 10, 26, 10]);
+  });
+
+  tearDownAll(() async {
+    await driver.close();
+  });
+}