From d537834b38c84a3815ebfa662272daa6732c480a Mon Sep 17 00:00:00 2001
From: Aubrey Anderson <aubrey@bigsur.com>
Date: Wed, 15 Apr 2020 18:25:02 -0700
Subject: [PATCH] Allow headers to be passed to the WebSocket connection for
 VMServiceFlutterDriver (#54698)

---
 .../flutter_driver/lib/src/driver/driver.dart  |  6 ++++++
 .../lib/src/driver/vmservice_driver.dart       | 17 +++++++++++------
 .../test/flutter_driver_test.dart              | 18 +++++++++++++++++-
 3 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/packages/flutter_driver/lib/src/driver/driver.dart b/packages/flutter_driver/lib/src/driver/driver.dart
index d9dfa7cb41..e06a7d569f 100644
--- a/packages/flutter_driver/lib/src/driver/driver.dart
+++ b/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,
     );
   }
 
diff --git a/packages/flutter_driver/lib/src/driver/vmservice_driver.dart b/packages/flutter_driver/lib/src/driver/vmservice_driver.dart
index 72063d12d0..52d3bcf92d 100644
--- a/packages/flutter_driver/lib/src/driver/vmservice_driver.dart
+++ b/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});
diff --git a/packages/flutter_driver/test/flutter_driver_test.dart b/packages/flutter_driver/test/flutter_driver_test.dart
index 18add5d4b5..18906b3663 100644
--- a/packages/flutter_driver/test/flutter_driver_test.dart
+++ b/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', () {