Add WebSocket cookie tests.

As it turned out WebSocket implementation had some problems with Cookie, this
CL adds unit tests.

BUG=None
R=rsleevi@chromium.org

Review URL: https://codereview.chromium.org/869073002

Cr-Commit-Position: refs/heads/master@{#315930}

7 files changed, 837 insertions, 197 deletions

diff --git a/net/net.gypi b/net/net.gypi
index 12d0a40..1098be4 100644
--- a/net/net.gypi
+++ b/net/net.gypi
@@ -1720,6 +1720,9 @@
       'websockets/websocket_handshake_challenge_test.cc',
       'websockets/websocket_handshake_stream_create_helper_test.cc',
       'websockets/websocket_inflater_test.cc',
+      'websockets/websocket_stream_cookie_test.cc',
+      'websockets/websocket_stream_create_test_base.cc',
+      'websockets/websocket_stream_create_test_base.h',
       'websockets/websocket_stream_test.cc',
       'websockets/websocket_test_util.cc',
       'websockets/websocket_test_util.h',
diff --git a/net/websockets/websocket_stream_cookie_test.cc b/net/websockets/websocket_stream_cookie_test.cc
new file mode 100644
index 0000000..9b819b5
--- /dev/null
+++ b/net/websockets/websocket_stream_cookie_test.cc
@@ -0,0 +1,503 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <string>
+
+#include "base/callback_forward.h"
+#include "base/memory/weak_ptr.h"
+#include "base/message_loop/message_loop.h"
+#include "base/run_loop.h"
+#include "base/strings/string_util.h"
+#include "base/strings/stringprintf.h"
+#include "base/thread_task_runner_handle.h"
+#include "net/cookies/cookie_store.h"
+#include "net/socket/socket_test_util.h"
+#include "net/websockets/websocket_stream_create_test_base.h"
+#include "net/websockets/websocket_test_util.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "url/gurl.h"
+
+namespace net {
+namespace {
+
+using ::testing::TestWithParam;
+using ::testing::ValuesIn;
+
+const char kNoCookieHeader[] = "";
+
+class TestBase : public WebSocketStreamCreateTestBase {
+ public:
+  void CreateAndConnect(const GURL& url,
+                        const std::string& origin,
+                        const std::string& cookie_header,
+                        const std::string& response_body) {
+    // We assume cookie_header ends with CRLF if not empty, as
+    // WebSocketStandardRequestWithCookies requires. Use AddCRLFIfNotEmpty
+    // in a call site.
+    CHECK(cookie_header.empty() || EndsWith(cookie_header, "\r\n", true));
+
+    url_request_context_host_.SetExpectations(
+        WebSocketStandardRequestWithCookies(url.path(), url.host(), origin,
+                                            cookie_header, std::string()),
+        response_body);
+    CreateAndConnectStream(url.spec(), NoSubProtocols(), origin, nullptr);
+  }
+
+  std::string AddCRLFIfNotEmpty(const std::string& s) {
+    return s.empty() ? s : s + "\r\n";
+  }
+};
+
+struct ClientUseCookieParameter {
+  // The URL for the WebSocket connection.
+  const char* const url;
+  // The URL for the previously set cookies.
+  const char* const cookie_url;
+  // The previously set cookies contents.
+  const char* const cookie_line;
+  // The Cookie: HTTP header expected to appear in the WS request. An empty
+  // string means there is no Cookie: header.
+  const char* const cookie_header;
+};
+
+class WebSocketStreamClientUseCookieTest
+    : public TestBase,
+      public TestWithParam<ClientUseCookieParameter> {
+ public:
+  ~WebSocketStreamClientUseCookieTest() override {
+    // Permit any endpoint locks to be released.
+    stream_request_.reset();
+    stream_.reset();
+    base::RunLoop().RunUntilIdle();
+  }
+
+  static void SetCookieHelperFunction(const base::Closure& task,
+                                      base::WeakPtr<bool> weak_is_called,
+                                      base::WeakPtr<bool> weak_result,
+                                      bool success) {
+    *weak_is_called = true;
+    *weak_result = success;
+    base::ThreadTaskRunnerHandle::Get()->PostTask(FROM_HERE, task);
+  }
+};
+
+struct ServerSetCookieParameter {
+  // The URL for the WebSocket connection.
+  const char* const url;
+  // The URL used to query cookies after the response received.
+  const char* const cookie_url;
+  // The cookies expected to appear for |cookie_url| inquiry.
+  const char* const cookie_line;
+  // The Set-Cookie: HTTP header attached to the response.
+  const char* const cookie_header;
+};
+
+class WebSocketStreamServerSetCookieTest
+    : public TestBase,
+      public TestWithParam<ServerSetCookieParameter> {
+ public:
+  ~WebSocketStreamServerSetCookieTest() override {
+    // Permit any endpoint locks to be released.
+    stream_request_.reset();
+    stream_.reset();
+    base::RunLoop().RunUntilIdle();
+  }
+
+  static void GetCookiesHelperFunction(const base::Closure& task,
+                                       base::WeakPtr<bool> weak_is_called,
+                                       base::WeakPtr<std::string> weak_result,
+                                       const std::string& cookies) {
+    *weak_is_called = true;
+    *weak_result = cookies;
+    base::ThreadTaskRunnerHandle::Get()->PostTask(FROM_HERE, task);
+  }
+};
+
+TEST_P(WebSocketStreamClientUseCookieTest, ClientUseCookie) {
+  // For wss tests.
+  ssl_data_.push_back(new SSLSocketDataProvider(ASYNC, OK));
+
+  CookieStore* store =
+      url_request_context_host_.GetURLRequestContext()->cookie_store();
+
+  const GURL url(GetParam().url);
+  const GURL cookie_url(GetParam().cookie_url);
+  const std::string origin("http://www.example.com");
+  const std::string cookie_line(GetParam().cookie_line);
+  const std::string cookie_header(AddCRLFIfNotEmpty(GetParam().cookie_header));
+
+  bool is_called = false;
+  bool set_cookie_result = false;
+  base::WeakPtrFactory<bool> weak_is_called(&is_called);
+  base::WeakPtrFactory<bool> weak_set_cookie_result(&set_cookie_result);
+
+  base::RunLoop run_loop;
+  store->SetCookieWithOptionsAsync(
+      cookie_url, cookie_line, CookieOptions(),
+      base::Bind(&SetCookieHelperFunction, run_loop.QuitClosure(),
+                 weak_is_called.GetWeakPtr(),
+                 weak_set_cookie_result.GetWeakPtr()));
+  run_loop.Run();
+  ASSERT_TRUE(is_called);
+  ASSERT_TRUE(set_cookie_result);
+
+  CreateAndConnect(url, origin, cookie_header, WebSocketStandardResponse(""));
+  WaitUntilConnectDone();
+  EXPECT_FALSE(has_failed());
+}
+
+TEST_P(WebSocketStreamServerSetCookieTest, ServerSetCookie) {
+  // For wss tests.
+  ssl_data_.push_back(new SSLSocketDataProvider(ASYNC, OK));
+
+  const GURL url(GetParam().url);
+  const GURL cookie_url(GetParam().cookie_url);
+  const std::string origin("http://www.example.com");
+  const std::string cookie_line(GetParam().cookie_line);
+  const std::string cookie_header(AddCRLFIfNotEmpty(GetParam().cookie_header));
+
+  const std::string response = base::StringPrintf(
+      "HTTP/1.1 101 Switching Protocols\r\n"
+      "Upgrade: websocket\r\n"
+      "Connection: Upgrade\r\n"
+      "%s"
+      "Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n"
+      "\r\n",
+      cookie_header.c_str());
+
+  CookieStore* store =
+      url_request_context_host_.GetURLRequestContext()->cookie_store();
+
+  CreateAndConnect(url, origin, "", response);
+  WaitUntilConnectDone();
+  EXPECT_FALSE(has_failed());
+
+  bool is_called = false;
+  std::string get_cookies_result;
+  base::WeakPtrFactory<bool> weak_is_called(&is_called);
+  base::WeakPtrFactory<std::string> weak_get_cookies_result(
+      &get_cookies_result);
+  base::RunLoop run_loop;
+  store->GetCookiesWithOptionsAsync(
+      cookie_url, CookieOptions(),
+      base::Bind(&GetCookiesHelperFunction, run_loop.QuitClosure(),
+                 weak_is_called.GetWeakPtr(),
+                 weak_get_cookies_result.GetWeakPtr()));
+  run_loop.Run();
+  EXPECT_TRUE(is_called);
+  EXPECT_EQ(cookie_line, get_cookies_result);
+}
+
+// Test parameters definitions follow...
+
+const ClientUseCookieParameter kClientUseCookieParameters[] = {
+    // Non-secure cookies for ws
+    {"ws://www.example.com",
+     "http://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "https://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "ws://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "wss://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    // Non-secure cookies for wss
+    {"wss://www.example.com",
+     "http://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "https://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "ws://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "wss://www.example.com",
+     "test-cookie",
+     "Cookie: test-cookie"},
+
+    // Secure-cookies for ws
+    {"ws://www.example.com",
+     "https://www.example.com",
+     "test-cookie; secure",
+     kNoCookieHeader},
+
+    {"ws://www.example.com",
+     "wss://www.example.com",
+     "test-cookie; secure",
+     kNoCookieHeader},
+
+    // Secure-cookies for wss
+    {"wss://www.example.com",
+     "https://www.example.com",
+     "test-cookie; secure",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "wss://www.example.com",
+     "test-cookie; secure",
+     "Cookie: test-cookie"},
+
+    // Non-secure cookies for ws (sharing domain)
+    {"ws://www.example.com",
+     "http://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "https://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "ws://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    // Non-secure cookies for wss (sharing domain)
+    {"wss://www.example.com",
+     "http://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "https://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "ws://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie; Domain=example.com",
+     "Cookie: test-cookie"},
+
+    // Secure-cookies for ws (sharing domain)
+    {"ws://www.example.com",
+     "https://www2.example.com",
+     "test-cookie; Domain=example.com; secure",
+     kNoCookieHeader},
+
+    {"ws://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie; Domain=example.com; secure",
+     kNoCookieHeader},
+
+    // Secure-cookies for wss (sharing domain)
+    {"wss://www.example.com",
+     "https://www2.example.com",
+     "test-cookie; Domain=example.com; secure",
+     "Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie; Domain=example.com; secure",
+     "Cookie: test-cookie"},
+
+    // Non-matching cookies for ws
+    {"ws://www.example.com",
+     "http://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+
+    {"ws://www.example.com",
+     "https://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+
+    {"ws://www.example.com",
+     "ws://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+
+    {"ws://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+
+    // Non-matching cookies for wss
+    {"wss://www.example.com",
+     "http://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+
+    {"wss://www.example.com",
+     "https://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+
+    {"wss://www.example.com",
+     "ws://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+
+    {"wss://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie",
+     kNoCookieHeader},
+};
+
+INSTANTIATE_TEST_CASE_P(WebSocketStreamClientUseCookieTest,
+                        WebSocketStreamClientUseCookieTest,
+                        ValuesIn(kClientUseCookieParameters));
+
+const ServerSetCookieParameter kServerSetCookieParameters[] = {
+    // Cookies coming from ws
+    {"ws://www.example.com",
+     "http://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "https://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "ws://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "wss://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    // Cookies coming from wss
+    {"wss://www.example.com",
+     "http://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "https://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "ws://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "wss://www.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie"},
+
+    // cookies coming from ws (sharing domain)
+    {"ws://www.example.com",
+     "http://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    {"ws://www.example.com",
+     "https://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    {"ws://www.example.com",
+     "ws://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    {"ws://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    // cookies coming from wss (sharing domain)
+    {"wss://www.example.com",
+     "http://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    {"wss://www.example.com",
+     "https://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    {"wss://www.example.com",
+     "ws://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    {"wss://www.example.com",
+     "wss://www2.example.com",
+     "test-cookie",
+     "Set-Cookie: test-cookie; Domain=example.com"},
+
+    // Non-matching cookies coming from ws
+    {"ws://www.example.com",
+     "http://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "https://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "ws://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+
+    {"ws://www.example.com",
+     "wss://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+
+    // Non-matching cookies coming from wss
+    {"wss://www.example.com",
+     "http://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "https://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "ws://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+
+    {"wss://www.example.com",
+     "wss://www2.example.com",
+     "",
+     "Set-Cookie: test-cookie"},
+};
+
+INSTANTIATE_TEST_CASE_P(WebSocketStreamServerSetCookieTest,
+                        WebSocketStreamServerSetCookieTest,
+                        ValuesIn(kServerSetCookieParameters));
+
+}  // namespace
+}  // namespace net
diff --git a/net/websockets/websocket_stream_create_test_base.cc b/net/websockets/websocket_stream_create_test_base.cc
new file mode 100644
index 0000000..900b1af
--- /dev/null
+++ b/net/websockets/websocket_stream_create_test_base.cc
@@ -0,0 +1,150 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "net/websockets/websocket_stream_create_test_base.h"
+
+#include "base/callback.h"
+#include "net/http/http_request_headers.h"
+#include "net/http/http_response_headers.h"
+#include "net/websockets/websocket_basic_handshake_stream.h"
+#include "net/websockets/websocket_handshake_request_info.h"
+#include "net/websockets/websocket_handshake_response_info.h"
+#include "net/websockets/websocket_handshake_stream_create_helper.h"
+#include "net/websockets/websocket_stream.h"
+#include "url/gurl.h"
+#include "url/origin.h"
+
+namespace net {
+
+using HeaderKeyValuePair = WebSocketStreamCreateTestBase::HeaderKeyValuePair;
+
+// A sub-class of WebSocketHandshakeStreamCreateHelper which always sets a
+// deterministic key to use in the WebSocket handshake.
+class DeterministicKeyWebSocketHandshakeStreamCreateHelper
+    : public WebSocketHandshakeStreamCreateHelper {
+ public:
+  DeterministicKeyWebSocketHandshakeStreamCreateHelper(
+      WebSocketStream::ConnectDelegate* connect_delegate,
+      const std::vector<std::string>& requested_subprotocols)
+      : WebSocketHandshakeStreamCreateHelper(connect_delegate,
+                                             requested_subprotocols) {}
+
+  void OnStreamCreated(WebSocketBasicHandshakeStream* stream) override {
+    stream->SetWebSocketKeyForTesting("dGhlIHNhbXBsZSBub25jZQ==");
+  }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(
+      DeterministicKeyWebSocketHandshakeStreamCreateHelper);
+};
+
+class WebSocketStreamCreateTestBase::TestConnectDelegate
+    : public WebSocketStream::ConnectDelegate {
+ public:
+  TestConnectDelegate(WebSocketStreamCreateTestBase* owner,
+                      const base::Closure& done_callback)
+      : owner_(owner), done_callback_(done_callback) {}
+
+  void OnSuccess(scoped_ptr<WebSocketStream> stream) override {
+    stream.swap(owner_->stream_);
+    done_callback_.Run();
+  }
+
+  void OnFailure(const std::string& message) override {
+    owner_->has_failed_ = true;
+    owner_->failure_message_ = message;
+    done_callback_.Run();
+  }
+
+  void OnStartOpeningHandshake(
+      scoped_ptr<WebSocketHandshakeRequestInfo> request) override {
+    // Can be called multiple times (in the case of HTTP auth). Last call
+    // wins.
+    owner_->request_info_ = request.Pass();
+  }
+
+  void OnFinishOpeningHandshake(
+      scoped_ptr<WebSocketHandshakeResponseInfo> response) override {
+    if (owner_->response_info_)
+      ADD_FAILURE();
+    owner_->response_info_ = response.Pass();
+  }
+
+  void OnSSLCertificateError(
+      scoped_ptr<WebSocketEventInterface::SSLErrorCallbacks>
+          ssl_error_callbacks,
+      const SSLInfo& ssl_info,
+      bool fatal) override {
+    owner_->ssl_error_callbacks_ = ssl_error_callbacks.Pass();
+    owner_->ssl_info_ = ssl_info;
+    owner_->ssl_fatal_ = fatal;
+  }
+
+ private:
+  WebSocketStreamCreateTestBase* owner_;
+  base::Closure done_callback_;
+  DISALLOW_COPY_AND_ASSIGN(TestConnectDelegate);
+};
+
+WebSocketStreamCreateTestBase::WebSocketStreamCreateTestBase()
+    : has_failed_(false), ssl_fatal_(false) {
+}
+
+WebSocketStreamCreateTestBase::~WebSocketStreamCreateTestBase() {
+}
+
+void WebSocketStreamCreateTestBase::CreateAndConnectStream(
+    const std::string& socket_url,
+    const std::vector<std::string>& sub_protocols,
+    const std::string& origin,
+    scoped_ptr<base::Timer> timer) {
+  for (size_t i = 0; i < ssl_data_.size(); ++i) {
+    scoped_ptr<SSLSocketDataProvider> ssl_data(ssl_data_[i]);
+    url_request_context_host_.AddSSLSocketDataProvider(ssl_data.Pass());
+  }
+  ssl_data_.weak_clear();
+  scoped_ptr<WebSocketStream::ConnectDelegate> connect_delegate(
+      new TestConnectDelegate(this, connect_run_loop_.QuitClosure()));
+  WebSocketStream::ConnectDelegate* delegate = connect_delegate.get();
+  scoped_ptr<WebSocketHandshakeStreamCreateHelper> create_helper(
+      new DeterministicKeyWebSocketHandshakeStreamCreateHelper(delegate,
+                                                               sub_protocols));
+  stream_request_ = CreateAndConnectStreamForTesting(
+      GURL(socket_url), create_helper.Pass(), url::Origin(origin),
+      url_request_context_host_.GetURLRequestContext(), BoundNetLog(),
+      connect_delegate.Pass(),
+      timer ? timer.Pass()
+            : scoped_ptr<base::Timer>(new base::Timer(false, false)));
+}
+
+std::vector<HeaderKeyValuePair>
+WebSocketStreamCreateTestBase::RequestHeadersToVector(
+    const HttpRequestHeaders& headers) {
+  HttpRequestHeaders::Iterator it(headers);
+  std::vector<HeaderKeyValuePair> result;
+  while (it.GetNext())
+    result.push_back(HeaderKeyValuePair(it.name(), it.value()));
+  return result;
+}
+
+std::vector<HeaderKeyValuePair>
+WebSocketStreamCreateTestBase::ResponseHeadersToVector(
+    const HttpResponseHeaders& headers) {
+  void* iter = NULL;
+  std::string name, value;
+  std::vector<HeaderKeyValuePair> result;
+  while (headers.EnumerateHeaderLines(&iter, &name, &value))
+    result.push_back(HeaderKeyValuePair(name, value));
+  return result;
+}
+
+void WebSocketStreamCreateTestBase::WaitUntilConnectDone() {
+  connect_run_loop_.Run();
+}
+
+std::vector<std::string> WebSocketStreamCreateTestBase::NoSubProtocols() {
+  return std::vector<std::string>();
+}
+
+}  // namespace net
diff --git a/net/websockets/websocket_stream_create_test_base.h b/net/websockets/websocket_stream_create_test_base.h
new file mode 100644
index 0000000..2388325
--- /dev/null
+++ b/net/websockets/websocket_stream_create_test_base.h
@@ -0,0 +1,87 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef NET_WEBSOCKETS_WEBSOCKET_STREAM_CREATE_TEST_BASE_H_
+#define NET_WEBSOCKETS_WEBSOCKET_STREAM_CREATE_TEST_BASE_H_
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "base/memory/scoped_ptr.h"
+#include "base/memory/scoped_vector.h"
+#include "base/run_loop.h"
+#include "base/timer/timer.h"
+#include "net/base/net_export.h"
+#include "net/socket/socket_test_util.h"
+#include "net/ssl/ssl_info.h"
+#include "net/websockets/websocket_event_interface.h"
+#include "net/websockets/websocket_test_util.h"
+
+namespace net {
+
+class HttpRequestHeaders;
+class HttpResponseHeaders;
+class WebSocketStream;
+class WebSocketStreamRequest;
+struct WebSocketHandshakeRequestInfo;
+struct WebSocketHandshakeResponseInfo;
+
+class WebSocketStreamCreateTestBase {
+ public:
+  using HeaderKeyValuePair = std::pair<std::string, std::string>;
+
+  WebSocketStreamCreateTestBase();
+  virtual ~WebSocketStreamCreateTestBase();
+
+  // A wrapper for CreateAndConnectStreamForTesting that knows about our default
+  // parameters.
+  void CreateAndConnectStream(const std::string& socket_url,
+                              const std::vector<std::string>& sub_protocols,
+                              const std::string& origin,
+                              scoped_ptr<base::Timer> timer);
+
+  static std::vector<HeaderKeyValuePair> RequestHeadersToVector(
+      const HttpRequestHeaders& headers);
+  static std::vector<HeaderKeyValuePair> ResponseHeadersToVector(
+      const HttpResponseHeaders& headers);
+
+  const std::string& failure_message() const { return failure_message_; }
+  bool has_failed() const { return has_failed_; }
+
+  // Runs |connect_run_loop_|. It will stop when the connection establishes or
+  // fails.
+  void WaitUntilConnectDone();
+
+  // A simple function to make the tests more readable.
+  std::vector<std::string> NoSubProtocols();
+
+ protected:
+  WebSocketTestURLRequestContextHost url_request_context_host_;
+  scoped_ptr<WebSocketStreamRequest> stream_request_;
+  // Only set if the connection succeeded.
+  scoped_ptr<WebSocketStream> stream_;
+  // Only set if the connection failed.
+  std::string failure_message_;
+  bool has_failed_;
+  scoped_ptr<WebSocketHandshakeRequestInfo> request_info_;
+  scoped_ptr<WebSocketHandshakeResponseInfo> response_info_;
+  scoped_ptr<WebSocketEventInterface::SSLErrorCallbacks> ssl_error_callbacks_;
+  SSLInfo ssl_info_;
+  bool ssl_fatal_;
+  ScopedVector<SSLSocketDataProvider> ssl_data_;
+
+  // This temporarily sets WebSocketEndpointLockManager unlock delay to zero
+  // during tests.
+  ScopedWebSocketEndpointZeroUnlockDelay zero_unlock_delay_;
+  base::RunLoop connect_run_loop_;
+
+ private:
+  class TestConnectDelegate;
+  DISALLOW_COPY_AND_ASSIGN(WebSocketStreamCreateTestBase);
+};
+
+}  // namespace net
+
+#endif  // NET_WEBSOCKETS_WEBSOCKET_STREAM_CREATE_TEST_BASE_H_
diff --git a/net/websockets/websocket_stream_test.cc b/net/websockets/websocket_stream_test.cc
index 8cfc0f1..56f2e64 100644
--- a/net/websockets/websocket_stream_test.cc
+++ b/net/websockets/websocket_stream_test.cc
@@ -29,9 +29,7 @@
 #include "net/url_request/url_request_test_util.h"
 #include "net/websockets/websocket_basic_handshake_stream.h"
 #include "net/websockets/websocket_frame.h"
-#include "net/websockets/websocket_handshake_request_info.h"
-#include "net/websockets/websocket_handshake_response_info.h"
-#include "net/websockets/websocket_handshake_stream_create_helper.h"
+#include "net/websockets/websocket_stream_create_test_base.h"
 #include "net/websockets/websocket_test_util.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "url/gurl.h"
@@ -40,25 +38,6 @@
 namespace net {
 namespace {
 
-typedef std::pair<std::string, std::string> HeaderKeyValuePair;
-
-std::vector<HeaderKeyValuePair> ToVector(const HttpRequestHeaders& headers) {
-  HttpRequestHeaders::Iterator it(headers);
-  std::vector<HeaderKeyValuePair> result;
-  while (it.GetNext())
-    result.push_back(HeaderKeyValuePair(it.name(), it.value()));
-  return result;
-}
-
-std::vector<HeaderKeyValuePair> ToVector(const HttpResponseHeaders& headers) {
-  void* iter = NULL;
-  std::string name, value;
-  std::vector<HeaderKeyValuePair> result;
-  while (headers.EnumerateHeaderLines(&iter, &name, &value))
-    result.push_back(HeaderKeyValuePair(name, value));
-  return result;
-}
-
 // Simple builder for a DeterministicSocketData object to save repetitive code.
 // It always sets the connect data to MockConnect(SYNCHRONOUS, OK), so it cannot
 // be used in tests where the connect fails. In practice, those tests never have
@@ -89,30 +68,14 @@ class MockWeakTimer : public base::MockTimer,
       : MockTimer(retain_user_task, is_repeating) {}
 };
 
-// A sub-class of WebSocketHandshakeStreamCreateHelper which always sets a
-// deterministic key to use in the WebSocket handshake.
-class DeterministicKeyWebSocketHandshakeStreamCreateHelper
-    : public WebSocketHandshakeStreamCreateHelper {
- public:
-  DeterministicKeyWebSocketHandshakeStreamCreateHelper(
-      WebSocketStream::ConnectDelegate* connect_delegate,
-      const std::vector<std::string>& requested_subprotocols)
-      : WebSocketHandshakeStreamCreateHelper(connect_delegate,
-                                             requested_subprotocols) {}
-
-  void OnStreamCreated(WebSocketBasicHandshakeStream* stream) override {
-    stream->SetWebSocketKeyForTesting("dGhlIHNhbXBsZSBub25jZQ==");
-  }
-};
-
-class WebSocketStreamCreateTest : public ::testing::Test {
+class WebSocketStreamCreateTest : public ::testing::Test,
+                                  public WebSocketStreamCreateTestBase {
  public:
-  WebSocketStreamCreateTest() : has_failed_(false), ssl_fatal_(false) {}
   ~WebSocketStreamCreateTest() override {
     // Permit any endpoint locks to be released.
     stream_request_.reset();
     stream_.reset();
-    RunUntilIdle();
+    base::RunLoop().RunUntilIdle();
   }
 
   void CreateAndConnectCustomResponse(
@@ -162,100 +125,6 @@ class WebSocketStreamCreateTest : public ::testing::Test {
   void AddRawExpectations(scoped_ptr<DeterministicSocketData> socket_data) {
     url_request_context_host_.AddRawExpectations(socket_data.Pass());
   }
-
-  // A wrapper for CreateAndConnectStreamForTesting that knows about our default
-  // parameters.
-  void CreateAndConnectStream(const std::string& socket_url,
-                              const std::vector<std::string>& sub_protocols,
-                              const std::string& origin,
-                              scoped_ptr<base::Timer> timer) {
-    for (size_t i = 0; i < ssl_data_.size(); ++i) {
-      scoped_ptr<SSLSocketDataProvider> ssl_data(ssl_data_[i]);
-      ssl_data_[i] = NULL;
-      url_request_context_host_.AddSSLSocketDataProvider(ssl_data.Pass());
-    }
-    ssl_data_.clear();
-    scoped_ptr<WebSocketStream::ConnectDelegate> connect_delegate(
-        new TestConnectDelegate(this));
-    WebSocketStream::ConnectDelegate* delegate = connect_delegate.get();
-    scoped_ptr<WebSocketHandshakeStreamCreateHelper> create_helper(
-        new DeterministicKeyWebSocketHandshakeStreamCreateHelper(
-            delegate, sub_protocols));
-    stream_request_ = ::net::CreateAndConnectStreamForTesting(
-        GURL(socket_url),
-        create_helper.Pass(),
-        url::Origin(origin),
-        url_request_context_host_.GetURLRequestContext(),
-        BoundNetLog(),
-        connect_delegate.Pass(),
-        timer ? timer.Pass() : scoped_ptr<base::Timer>(
-            new base::Timer(false, false)));
-  }
-
-  static void RunUntilIdle() { base::RunLoop().RunUntilIdle(); }
-
-  // A simple function to make the tests more readable. Creates an empty vector.
-  static std::vector<std::string> NoSubProtocols() {
-    return std::vector<std::string>();
-  }
-
-  const std::string& failure_message() const { return failure_message_; }
-  bool has_failed() const { return has_failed_; }
-
-  class TestConnectDelegate : public WebSocketStream::ConnectDelegate {
-   public:
-    explicit TestConnectDelegate(WebSocketStreamCreateTest* owner)
-        : owner_(owner) {}
-
-    void OnSuccess(scoped_ptr<WebSocketStream> stream) override {
-      stream.swap(owner_->stream_);
-    }
-
-    void OnFailure(const std::string& message) override {
-      owner_->has_failed_ = true;
-      owner_->failure_message_ = message;
-    }
-
-    void OnStartOpeningHandshake(
-        scoped_ptr<WebSocketHandshakeRequestInfo> request) override {
-      // Can be called multiple times (in the case of HTTP auth). Last call
-      // wins.
-      owner_->request_info_ = request.Pass();
-    }
-    void OnFinishOpeningHandshake(
-        scoped_ptr<WebSocketHandshakeResponseInfo> response) override {
-      if (owner_->response_info_)
-        ADD_FAILURE();
-      owner_->response_info_ = response.Pass();
-    }
-    void OnSSLCertificateError(
-        scoped_ptr<WebSocketEventInterface::SSLErrorCallbacks>
-            ssl_error_callbacks,
-        const SSLInfo& ssl_info,
-        bool fatal) override {
-      owner_->ssl_error_callbacks_ = ssl_error_callbacks.Pass();
-      owner_->ssl_info_ = ssl_info;
-      owner_->ssl_fatal_ = fatal;
-    }
-
-   private:
-    WebSocketStreamCreateTest* owner_;
-  };
-
-  WebSocketTestURLRequestContextHost url_request_context_host_;
-  scoped_ptr<WebSocketStreamRequest> stream_request_;
-  // Only set if the connection succeeded.
-  scoped_ptr<WebSocketStream> stream_;
-  // Only set if the connection failed.
-  std::string failure_message_;
-  bool has_failed_;
-  scoped_ptr<WebSocketHandshakeRequestInfo> request_info_;
-  scoped_ptr<WebSocketHandshakeResponseInfo> response_info_;
-  scoped_ptr<WebSocketEventInterface::SSLErrorCallbacks> ssl_error_callbacks_;
-  SSLInfo ssl_info_;
-  bool ssl_fatal_;
-  ScopedVector<SSLSocketDataProvider> ssl_data_;
-  ScopedWebSocketEndpointZeroUnlockDelay zero_unlock_delay_;
 };
 
 // There are enough tests of the Sec-WebSocket-Extensions header that they
@@ -271,7 +140,7 @@ class WebSocketStreamCreateExtensionTest : public WebSocketStreamCreateTest {
         "ws://localhost/testing_path", "localhost", "/testing_path",
         NoSubProtocols(), "http://localhost", "",
         "Sec-WebSocket-Extensions: " + extensions_header_value + "\r\n");
-    RunUntilIdle();
+    WaitUntilConnectDone();
   }
 };
 
@@ -430,7 +299,7 @@ TEST_F(WebSocketStreamCreateTest, SimpleSuccess) {
                            NoSubProtocols(), "http://localhost", "", "");
   EXPECT_FALSE(request_info_);
   EXPECT_FALSE(response_info_);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
   EXPECT_TRUE(request_info_);
@@ -453,12 +322,12 @@ TEST_F(WebSocketStreamCreateTest, HandshakeInfo) {
                                  kResponse);
   EXPECT_FALSE(request_info_);
   EXPECT_FALSE(response_info_);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(stream_);
   ASSERT_TRUE(request_info_);
   ASSERT_TRUE(response_info_);
   std::vector<HeaderKeyValuePair> request_headers =
-      ToVector(request_info_->headers);
+      RequestHeadersToVector(request_info_->headers);
   // We examine the contents of request_info_ and response_info_
   // mainly only in this test case.
   EXPECT_EQ(GURL("ws://localhost/"), request_info_->url);
@@ -487,7 +356,7 @@ TEST_F(WebSocketStreamCreateTest, HandshakeInfo) {
             request_headers[11]);
 
   std::vector<HeaderKeyValuePair> response_headers =
-      ToVector(*response_info_->headers.get());
+      ResponseHeadersToVector(*response_info_->headers.get());
   ASSERT_EQ(6u, response_headers.size());
   // Sort the headers for ease of verification.
   std::sort(response_headers.begin(), response_headers.end());
@@ -513,7 +382,7 @@ TEST_F(WebSocketStreamCreateTest, PathIsUsed) {
   CreateAndConnectStandard("ws://localhost/testing_path", "localhost",
                            "/testing_path", NoSubProtocols(),
                            "http://localhost", "", "");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
 }
@@ -523,7 +392,7 @@ TEST_F(WebSocketStreamCreateTest, OriginIsUsed) {
   CreateAndConnectStandard("ws://localhost/testing_path", "localhost",
                            "/testing_path", NoSubProtocols(),
                            "http://google.com", "", "");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
 }
@@ -538,7 +407,7 @@ TEST_F(WebSocketStreamCreateTest, SubProtocolIsUsed) {
                            "Sec-WebSocket-Protocol: chatv11.chromium.org, "
                            "chatv20.chromium.org\r\n",
                            "Sec-WebSocket-Protocol: chatv20.chromium.org\r\n");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(stream_);
   EXPECT_FALSE(has_failed());
   EXPECT_EQ("chatv20.chromium.org", stream_->GetSubProtocol());
@@ -550,7 +419,7 @@ TEST_F(WebSocketStreamCreateTest, UnsolicitedSubProtocol) {
                            "/testing_path", NoSubProtocols(),
                            "http://google.com", "",
                            "Sec-WebSocket-Protocol: chatv20.chromium.org\r\n");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(stream_);
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
@@ -566,7 +435,7 @@ TEST_F(WebSocketStreamCreateTest, UnacceptedSubProtocol) {
   CreateAndConnectStandard("ws://localhost/testing_path", "localhost",
                            "/testing_path", sub_protocols, "http://localhost",
                            "Sec-WebSocket-Protocol: chat.example.com\r\n", "");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(stream_);
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
@@ -586,7 +455,7 @@ TEST_F(WebSocketStreamCreateTest, MultipleSubProtocolsInResponse) {
                            "chatv20.chromium.org\r\n",
                            "Sec-WebSocket-Protocol: chatv11.chromium.org, "
                            "chatv20.chromium.org\r\n");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(stream_);
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
@@ -605,7 +474,7 @@ TEST_F(WebSocketStreamCreateTest, UnmatchedSubProtocolInResponse) {
                            "Sec-WebSocket-Protocol: chatv11.chromium.org, "
                            "chatv20.chromium.org\r\n",
                            "Sec-WebSocket-Protocol: chatv21.chromium.org\r\n");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(stream_);
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
@@ -643,7 +512,7 @@ TEST_F(WebSocketStreamCreateExtensionTest, PerMessageDeflateInflates) {
               "\xc1\x07"  // WebSocket header (FIN + RSV1, Text payload 7 bytes)
               "\xf2\x48\xcd\xc9\xc9\x07\x00",  // "Hello" DEFLATE compressed
               9));
-  RunUntilIdle();
+  WaitUntilConnectDone();
 
   ASSERT_TRUE(stream_);
   ScopedVector<WebSocketFrame> frames;
@@ -820,7 +689,7 @@ TEST_F(WebSocketStreamCreateTest, DoubleAccept) {
   CreateAndConnectStandard(
       "ws://localhost/", "localhost", "/", NoSubProtocols(), "http://localhost",
       "", "Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(stream_);
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
@@ -840,7 +709,7 @@ TEST_F(WebSocketStreamCreateTest, InvalidStatusCode) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kInvalidStatusCodeResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: Unexpected response code: 200",
             failure_message());
@@ -860,7 +729,7 @@ TEST_F(WebSocketStreamCreateTest, RedirectsRejected) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kRedirectResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: Unexpected response code: 302",
             failure_message());
@@ -881,7 +750,7 @@ TEST_F(WebSocketStreamCreateTest, MalformedResponse) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kMalformedResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: Invalid status line",
             failure_message());
@@ -897,7 +766,7 @@ TEST_F(WebSocketStreamCreateTest, MissingUpgradeHeader) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kMissingUpgradeResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: 'Upgrade' header is missing",
             failure_message());
@@ -908,7 +777,7 @@ TEST_F(WebSocketStreamCreateTest, DoubleUpgradeHeader) {
   CreateAndConnectStandard("ws://localhost/", "localhost", "/",
                            NoSubProtocols(), "http://localhost", "",
                            "Upgrade: HTTP/2.0\r\n");
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
             "'Upgrade' header must not appear more than once in a response",
@@ -926,7 +795,7 @@ TEST_F(WebSocketStreamCreateTest, IncorrectUpgradeHeader) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kMissingUpgradeResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
             "'Upgrade' header value is not 'WebSocket': hogefuga",
@@ -943,7 +812,7 @@ TEST_F(WebSocketStreamCreateTest, MissingConnectionHeader) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kMissingConnectionResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
             "'Connection' header is missing",
@@ -961,7 +830,7 @@ TEST_F(WebSocketStreamCreateTest, IncorrectConnectionHeader) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kMissingConnectionResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
             "'Connection' header value must contain 'Upgrade'",
@@ -979,7 +848,7 @@ TEST_F(WebSocketStreamCreateTest, AdditionalTokenInConnectionHeader) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kAdditionalConnectionTokenResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
 }
@@ -994,7 +863,7 @@ TEST_F(WebSocketStreamCreateTest, MissingSecWebSocketAccept) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kMissingAcceptResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
             "'Sec-WebSocket-Accept' header is missing",
@@ -1012,7 +881,7 @@ TEST_F(WebSocketStreamCreateTest, WrongSecWebSocketAccept) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kIncorrectAcceptResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error during WebSocket handshake: "
             "Incorrect 'Sec-WebSocket-Accept' header value",
@@ -1024,7 +893,8 @@ TEST_F(WebSocketStreamCreateTest, Cancellation) {
   CreateAndConnectStandard("ws://localhost/", "localhost", "/",
                            NoSubProtocols(), "http://localhost", "", "");
   stream_request_.reset();
-  RunUntilIdle();
+  // WaitUntilConnectDone doesn't work in this case.
+  base::RunLoop().RunUntilIdle();
   EXPECT_FALSE(has_failed());
   EXPECT_FALSE(stream_);
   EXPECT_FALSE(request_info_);
@@ -1038,7 +908,7 @@ TEST_F(WebSocketStreamCreateTest, ConnectionFailure) {
       MockConnect(SYNCHRONOUS, ERR_CONNECTION_REFUSED));
   CreateAndConnectRawExpectations("ws://localhost/", NoSubProtocols(),
                                   "http://localhost", socket_data.Pass());
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error in connection establishment: net::ERR_CONNECTION_REFUSED",
             failure_message());
@@ -1053,7 +923,7 @@ TEST_F(WebSocketStreamCreateTest, ConnectionTimeout) {
       MockConnect(ASYNC, ERR_CONNECTION_TIMED_OUT));
   CreateAndConnectRawExpectations("ws://localhost/", NoSubProtocols(),
                                   "http://localhost", socket_data.Pass());
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error in connection establishment: net::ERR_CONNECTION_TIMED_OUT",
             failure_message());
@@ -1075,7 +945,7 @@ TEST_F(WebSocketStreamCreateTest, HandshakeTimeout) {
   EXPECT_TRUE(weak_timer->IsRunning());
 
   weak_timer->Fire();
-  RunUntilIdle();
+  WaitUntilConnectDone();
 
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("WebSocket opening handshake timed out", failure_message());
@@ -1094,7 +964,7 @@ TEST_F(WebSocketStreamCreateTest, HandshakeTimerOnSuccess) {
   ASSERT_TRUE(weak_timer);
   EXPECT_TRUE(weak_timer->IsRunning());
 
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
   ASSERT_TRUE(weak_timer);
@@ -1116,7 +986,7 @@ TEST_F(WebSocketStreamCreateTest, HandshakeTimerOnFailure) {
   ASSERT_TRUE(weak_timer.get());
   EXPECT_TRUE(weak_timer->IsRunning());
 
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Error in connection establishment: net::ERR_CONNECTION_REFUSED",
             failure_message());
@@ -1133,7 +1003,8 @@ TEST_F(WebSocketStreamCreateTest, CancellationDuringConnect) {
                                   "http://localhost",
                                   socket_data.Pass());
   stream_request_.reset();
-  RunUntilIdle();
+  // WaitUntilConnectDone doesn't work in this case.
+  base::RunLoop().RunUntilIdle();
   EXPECT_FALSE(has_failed());
   EXPECT_FALSE(stream_);
 }
@@ -1154,7 +1025,8 @@ TEST_F(WebSocketStreamCreateTest, CancellationDuringWrite) {
                                   make_scoped_ptr(socket_data));
   socket_data->Run();
   stream_request_.reset();
-  RunUntilIdle();
+  // WaitUntilConnectDone doesn't work in this case.
+  base::RunLoop().RunUntilIdle();
   EXPECT_FALSE(has_failed());
   EXPECT_FALSE(stream_);
   EXPECT_TRUE(request_info_);
@@ -1179,7 +1051,8 @@ TEST_F(WebSocketStreamCreateTest, CancellationDuringRead) {
                                   socket_data.Pass());
   socket_data_raw_ptr->Run();
   stream_request_.reset();
-  RunUntilIdle();
+  // WaitUntilConnectDone doesn't work in this case.
+  base::RunLoop().RunUntilIdle();
   EXPECT_FALSE(has_failed());
   EXPECT_FALSE(stream_);
   EXPECT_TRUE(request_info_);
@@ -1199,7 +1072,7 @@ TEST_F(WebSocketStreamCreateTest, VeryLargeResponseHeaders) {
   CreateAndConnectStandard("ws://localhost/", "localhost", "/",
                            NoSubProtocols(), "http://localhost", "",
                            set_cookie_headers);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_FALSE(response_info_);
 }
@@ -1238,12 +1111,13 @@ TEST_F(WebSocketStreamCreateTest, SelfSignedCertificateFailure) {
                                   NoSubProtocols(),
                                   "http://localhost",
                                   raw_socket_data.Pass());
-  RunUntilIdle();
+  // WaitUntilConnectDone doesn't work in this case.
+  base::RunLoop().RunUntilIdle();
   EXPECT_FALSE(has_failed());
   ASSERT_TRUE(ssl_error_callbacks_);
   ssl_error_callbacks_->CancelSSLRequest(ERR_CERT_AUTHORITY_INVALID,
                                          &ssl_info_);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
 }
 
@@ -1259,10 +1133,11 @@ TEST_F(WebSocketStreamCreateTest, SelfSignedCertificateSuccess) {
   url_request_context_host_.AddRawExpectations(BuildNullSocketData());
   CreateAndConnectStandard("wss://localhost/", "localhost", "/",
                            NoSubProtocols(), "http://localhost", "", "");
-  RunUntilIdle();
+  // WaitUntilConnectDone doesn't work in this case.
+  base::RunLoop().RunUntilIdle();
   ASSERT_TRUE(ssl_error_callbacks_);
   ssl_error_callbacks_->ContinueSSLRequest();
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
 }
@@ -1273,7 +1148,7 @@ TEST_F(WebSocketStreamCreateBasicAuthTest, FailureNoCredentials) {
   CreateAndConnectCustomResponse("ws://localhost/", "localhost", "/",
                                  NoSubProtocols(), "http://localhost", "",
                                  kUnauthorizedResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("HTTP Authentication failed; no valid credentials available",
             failure_message());
@@ -1284,7 +1159,7 @@ TEST_F(WebSocketStreamCreateBasicAuthTest, SuccessPasswordInUrl) {
   CreateAndConnectAuthHandshake("ws://foo:bar@localhost/",
                                 "Zm9vOmJhcg==",
                                 WebSocketStandardResponse(std::string()));
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
   ASSERT_TRUE(response_info_);
@@ -1294,7 +1169,7 @@ TEST_F(WebSocketStreamCreateBasicAuthTest, SuccessPasswordInUrl) {
 TEST_F(WebSocketStreamCreateBasicAuthTest, FailureIncorrectPasswordInUrl) {
   CreateAndConnectAuthHandshake(
       "ws://foo:baz@localhost/", "Zm9vOmJheg==", kUnauthorizedResponse);
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_TRUE(response_info_);
 }
@@ -1311,7 +1186,7 @@ TEST_F(WebSocketStreamCreateDigestAuthTest, DigestPasswordInUrl) {
       "http://localhost",
       helper_.BuildSocketData2(kAuthorizedRequest,
                                WebSocketStandardResponse(std::string())));
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_FALSE(has_failed());
   EXPECT_TRUE(stream_);
   ASSERT_TRUE(response_info_);
@@ -1348,7 +1223,7 @@ TEST_F(WebSocketStreamCreateUMATest, Connected) {
     creation.CreateAndConnectStandard("ws://localhost/", "localhost", "/",
                                       creation.NoSubProtocols(),
                                       "http://localhost", "", "");
-    creation.RunUntilIdle();
+    creation.WaitUntilConnectDone();
   }
 
   scoped_ptr<base::HistogramSamples> samples(GetSamples(name));
@@ -1376,7 +1251,7 @@ TEST_F(WebSocketStreamCreateUMATest, Failed) {
     creation.CreateAndConnectCustomResponse(
         "ws://localhost/", "localhost", "/", creation.NoSubProtocols(),
         "http://localhost", "", kInvalidStatusCodeResponse);
-    creation.RunUntilIdle();
+    creation.WaitUntilConnectDone();
   }
 
   scoped_ptr<base::HistogramSamples> samples(GetSamples(name));
@@ -1409,7 +1284,7 @@ TEST_F(WebSocketStreamCreateTest, HandleErrConnectionClosed) {
   socket_data->set_connect_data(MockConnect(SYNCHRONOUS, OK));
   CreateAndConnectRawExpectations("ws://localhost/", NoSubProtocols(),
                                   "http://localhost", socket_data.Pass());
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
 }
 
@@ -1435,7 +1310,7 @@ TEST_F(WebSocketStreamCreateTest, HandleErrTunnelConnectionFailed) {
   url_request_context_host_.SetProxyConfig("https=proxy:8000");
   CreateAndConnectRawExpectations("ws://localhost/", NoSubProtocols(),
                                   "http://localhost", socket_data.Pass());
-  RunUntilIdle();
+  WaitUntilConnectDone();
   EXPECT_TRUE(has_failed());
   EXPECT_EQ("Establishing a tunnel via proxy server failed.",
             failure_message());
diff --git a/net/websockets/websocket_test_util.cc b/net/websockets/websocket_test_util.cc
index 0b0b8c0..3f963bf 100644
--- a/net/websockets/websocket_test_util.cc
+++ b/net/websockets/websocket_test_util.cc
@@ -37,6 +37,16 @@ std::string WebSocketStandardRequest(const std::string& path,
                                      const std::string& host,
                                      const std::string& origin,
                                      const std::string& extra_headers) {
+  return WebSocketStandardRequestWithCookies(path, host, origin, std::string(),
+                                             extra_headers);
+}
+
+std::string WebSocketStandardRequestWithCookies(
+    const std::string& path,
+    const std::string& host,
+    const std::string& origin,
+    const std::string& cookies,
+    const std::string& extra_headers) {
   // Unrelated changes in net/http may change the order and default-values of
   // HTTP headers, causing WebSocket tests to fail. It is safe to update this
   // string in that case.
@@ -52,10 +62,12 @@ std::string WebSocketStandardRequest(const std::string& path,
       "User-Agent:\r\n"
       "Accept-Encoding: gzip, deflate\r\n"
       "Accept-Language: en-us,fr\r\n"
+      "%s"
       "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
       "Sec-WebSocket-Extensions: permessage-deflate; client_max_window_bits\r\n"
       "%s\r\n",
-      path.c_str(), host.c_str(), origin.c_str(), extra_headers.c_str());
+      path.c_str(), host.c_str(), origin.c_str(), cookies.c_str(),
+      extra_headers.c_str());
 }
 
 std::string WebSocketStandardResponse(const std::string& extra_headers) {
diff --git a/net/websockets/websocket_test_util.h b/net/websockets/websocket_test_util.h
index f9bf961..0230f5c9 100644
--- a/net/websockets/websocket_test_util.h
+++ b/net/websockets/websocket_test_util.h
@@ -45,27 +45,37 @@ class LinearCongruentialGenerator {
 // use only. The differences are the use of a |create_helper| argument in place
 // of |requested_subprotocols| and taking |timer| as the handshake timeout
 // timer. Implemented in websocket_stream.cc.
-NET_EXPORT_PRIVATE extern scoped_ptr<WebSocketStreamRequest>
-    CreateAndConnectStreamForTesting(
-        const GURL& socket_url,
-        scoped_ptr<WebSocketHandshakeStreamCreateHelper> create_helper,
-        const url::Origin& origin,
-        URLRequestContext* url_request_context,
-        const BoundNetLog& net_log,
-        scoped_ptr<WebSocketStream::ConnectDelegate> connect_delegate,
-        scoped_ptr<base::Timer> timer);
+NET_EXPORT_PRIVATE scoped_ptr<WebSocketStreamRequest>
+CreateAndConnectStreamForTesting(
+    const GURL& socket_url,
+    scoped_ptr<WebSocketHandshakeStreamCreateHelper> create_helper,
+    const url::Origin& origin,
+    URLRequestContext* url_request_context,
+    const BoundNetLog& net_log,
+    scoped_ptr<WebSocketStream::ConnectDelegate> connect_delegate,
+    scoped_ptr<base::Timer> timer);
 
 // Generates a standard WebSocket handshake request. The challenge key used is
 // "dGhlIHNhbXBsZSBub25jZQ==". Each header in |extra_headers| must be terminated
 // with "\r\n".
-extern std::string WebSocketStandardRequest(const std::string& path,
-                                            const std::string& host,
-                                            const std::string& origin,
-                                            const std::string& extra_headers);
+std::string WebSocketStandardRequest(const std::string& path,
+                                     const std::string& host,
+                                     const std::string& origin,
+                                     const std::string& extra_headers);
+
+// Generates a standard WebSocket handshake request. The challenge key used is
+// "dGhlIHNhbXBsZSBub25jZQ==". |cookies| must be empty or terminated with
+// "\r\n". Each header in |extra_headers| must be terminated with "\r\n".
+std::string WebSocketStandardRequestWithCookies(
+    const std::string& path,
+    const std::string& host,
+    const std::string& origin,
+    const std::string& cookies,
+    const std::string& extra_headers);
 
 // A response with the appropriate accept header to match the above challenge
 // key. Each header in |extra_headers| must be terminated with "\r\n".
-extern std::string WebSocketStandardResponse(const std::string& extra_headers);
+std::string WebSocketStandardResponse(const std::string& extra_headers);
 
 // This class provides a convenient way to construct a
 // DeterministicMockClientSocketFactory for WebSocket tests.