Web Socket API for talking to Editor.

sdk/lib/_internal/pub/test/serve/utils.dart

Index: sdk/lib/_internal/pub/test/serve/utils.dart
diff --git a/sdk/lib/_internal/pub/test/serve/utils.dart b/sdk/lib/_internal/pub/test/serve/utils.dart
index 93ffbff2ff5b3ead988d3da803b5a84d98570a61..6ff10c0b74dad34081bc498c9293eba92851bfa2 100644
--- a/sdk/lib/_internal/pub/test/serve/utils.dart
+++ b/sdk/lib/_internal/pub/test/serve/utils.dart
@@ -5,6 +5,8 @@
 library pub_tests;
 
 import 'dart:async';
+import 'dart:convert';
+import 'dart:io';
 
 import 'package:http/http.dart' as http;
 import 'package:scheduled_test/scheduled_process.dart';
@@ -18,6 +20,11 @@ ScheduledProcess _pubServer;
 /// The ephemeral port assigned to the running server.
 int _port;
 
+/// The web socket connection to the running pub process, or `null` if no
+/// connection has been made.
+WebSocket _webSocket;
+Stream _webSocketBroadcastStream;
+
 /// The code for a transformer that renames ".txt" files to ".out" and adds a
 /// ".out" suffix.
 const REWRITE_TRANSFORMER = """
@@ -107,6 +114,14 @@ ScheduledProcess startPubServe({bool shouldGetFirst: false,
 
   _pubServer = startPub(args: args);
 
+  currentSchedule.onComplete.schedule(() {
+    if (_webSocket != null) {
+      _webSocket.close();
+      _webSocket = null;
+      _webSocketBroadcastStream = null;
+    }
+  });
+
   if (shouldGetFirst) {
     expect(_pubServer.nextLine(),
         completion(anyOf(
@@ -183,3 +198,36 @@ void waitForBuildSuccess() {
 
   schedule(nextLine);
 }
+
+/// Schedules opening a web socket connection to the currently running pub
+/// serve.
+Future _ensureWebSocket() {
+  // Use the existing one if already connected.
+  if (_webSocket != null) return new Future.value();
+
+  // Server should already be running.
+  assert(_pubServer != null);
+  assert(_port != null);
+
+  return WebSocket.connect("ws://127.0.0.1:$_port").then((socket) {
+    _webSocket = socket;
+    // TODO(rnystrom): Works around #13913.
+    _webSocketBroadcastStream = _webSocket.asBroadcastStream();
+  });
+}
+
+/// Sends [request] (an arbitrary JSON object) to the running pub serve's web
+/// socket connection, waits for a reply, then verifies that the reply matches
+/// [expectation].
+///
+/// If [encodeRequest] is `false`, then [request] will be sent as-is over the
+/// socket. It omitted, request is JSON encoded to a string first.
+void webSocketShouldReply(request, expectation, {bool encodeRequest: true}) {
+  schedule(() => _ensureWebSocket().then((_) {
+    if (encodeRequest) request = JSON.encode(request);
+    _webSocket.add(request);
+    return _webSocketBroadcastStream.first.then((value) {
+      expect(JSON.decode(value), expectation);
+    });
+  }), "send $request to web socket and expect reply that $expectation");
+}