Allow headers to be passed to the WebSocket connection for VMServiceFlutterDriver (#54698)

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

@@ -127,6 +127,10 @@ abstract class FlutterDriver {
   /// `isolateNumber` is set, as this is already enough information to connect
   /// to an isolate.
   ///
+  /// `headers` optionally specifies HTTP headers to be included in the
+  /// [WebSocket] connection. This is only used for [VMServiceFlutterDriver]
+  /// connections.
+  ///
   /// `browser` specifies which FlutterDriver implementation to use. If not
   /// speicifed or set to false, [VMServiceFlutterDriver] implementation
   /// will be used. Otherwise, [WebFlutterDriver] implementation will be used.
@@ -141,6 +145,7 @@ abstract class FlutterDriver {
     int isolateNumber,
     Pattern fuchsiaModuleTarget,
     Duration timeout,
+    Map<String, dynamic> headers,
   }) async {
     if (Platform.environment['FLUTTER_WEB_TEST'] != null) {
       return WebFlutterDriver.connectWeb(hostUrl: dartVmServiceUrl, timeout: timeout);
@@ -151,6 +156,7 @@ abstract class FlutterDriver {
               logCommunicationToFile: logCommunicationToFile,
               isolateNumber: isolateNumber,
               fuchsiaModuleTarget: fuchsiaModuleTarget,
+              headers: headers,
     );
   }
 

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

@@ -48,6 +48,7 @@ class VMServiceFlutterDriver extends FlutterDriver {
     bool logCommunicationToFile = true,
     int isolateNumber,
     Pattern fuchsiaModuleTarget,
+    Map<String, dynamic> headers,
   }) async {
     // If running on a Fuchsia device, connect to the first isolate whose name
     // matches FUCHSIA_MODULE_TARGET.
@@ -92,7 +93,7 @@ class VMServiceFlutterDriver extends FlutterDriver {
     // Connect to Dart VM services
     _log('Connecting to Flutter application at $dartVmServiceUrl');
     final VMServiceClientConnection connection =
-    await vmServiceConnectFunction(dartVmServiceUrl);
+    await vmServiceConnectFunction(dartVmServiceUrl, headers: headers);
     final VMServiceClient client = connection.client;
     final VM vm = await client.getVM();
     final VMIsolateRef isolateRef = isolateNumber ==
@@ -564,15 +565,16 @@ void _checkCloseCode(WebSocket ws) {
 
 /// Waits for a real Dart VM service to become available, then connects using
 /// the [VMServiceClient].
-Future<VMServiceClientConnection> _waitAndConnect(String url) async {
+Future<VMServiceClientConnection> _waitAndConnect(
+    String url, {Map<String, dynamic> headers}) async {
   final String webSocketUrl = _getWebSocketUrl(url);
   int attempts = 0;
   while (true) {
     WebSocket ws1;
     WebSocket ws2;
     try {
-      ws1 = await WebSocket.connect(webSocketUrl);
-      ws2 = await WebSocket.connect(webSocketUrl);
+      ws1 = await WebSocket.connect(webSocketUrl, headers: headers);
+      ws2 = await WebSocket.connect(webSocketUrl, headers: headers);
 
       ws1.done.whenComplete(() => _checkCloseCode(ws1));
       ws2.done.whenComplete(() => _checkCloseCode(ws2));
@@ -650,5 +652,8 @@ class VMServiceClientConnection {
   final rpc.Peer peer;
 }
 
-/// A function that connects to a Dart VM service given the [url].
-typedef VMServiceConnectFunction = Future<VMServiceClientConnection> Function(String url);
+/// A function that connects to a Dart VM service
+/// with [headers] given the [url].
+typedef VMServiceConnectFunction =
+  Future<VMServiceClientConnection> Function(
+    String url, {Map<String, dynamic> headers});

packages/flutter_driver/test/flutter_driver_test.dart

@@ -51,7 +51,7 @@ void main() {
       when(mockIsolate.loadRunnable()).thenAnswer((_) => Future<MockIsolate>.value(mockIsolate));
       when(mockIsolate.invokeExtension(any, any)).thenAnswer(
           (Invocation invocation) => makeMockResponse(<String, dynamic>{'status': 'ok'}));
-      vmServiceConnectFunction = (String url) {
+      vmServiceConnectFunction = (String url, {Map<String, dynamic> headers}) {
         return Future<VMServiceClientConnection>.value(
           VMServiceClientConnection(mockClient, mockPeer)
         );
@@ -116,6 +116,22 @@ void main() {
       expect(driver, isNotNull);
       expectLogContains('Isolate is not paused. Assuming application is ready.');
     });
+
+    test('connects with headers', () async {
+      Map<String, dynamic> actualHeaders;
+      vmServiceConnectFunction = (String url, {Map<String, dynamic> headers}) {
+        actualHeaders = headers;
+        return Future<VMServiceClientConnection>.value(
+          VMServiceClientConnection(mockClient, mockPeer)
+        );
+      };
+
+      final Map<String, String> expectedHeaders = <String, String>{'header-key': 'header-value'};
+      final FlutterDriver driver = await FlutterDriver.connect(
+        dartVmServiceUrl: '', headers: expectedHeaders);
+      expect(driver, isNotNull);
+      expect(actualHeaders, equals(expectedHeaders));
+    });
   });
 
   group('VMServiceFlutterDriver', () {