Refactor CastSocket code for the following:

chrome/browser/extensions/api/cast_channel/cast_socket_unittest.cc

Index: chrome/browser/extensions/api/cast_channel/cast_socket_unittest.cc
===================================================================
--- chrome/browser/extensions/api/cast_channel/cast_socket_unittest.cc	(revision 236353)
+++ chrome/browser/extensions/api/cast_channel/cast_socket_unittest.cc	(working copy)
@@ -4,6 +4,9 @@
 
 #include "chrome/browser/extensions/api/cast_channel/cast_socket.h"
 
+#include "base/message_loop/message_loop.h"
+#include "base/run_loop.h"
+#include "base/strings/string_number_conversions.h"
 #include "chrome/browser/extensions/api/cast_channel/cast_channel.pb.h"
 #include "chrome/browser/extensions/api/cast_channel/cast_message_util.h"
 #include "net/base/address_list.h"
@@ -24,6 +27,15 @@
 using ::testing::Return;
 using ::testing::SaveArg;
 
+namespace {
+const char* kTestData[4] = {
+    "Hello, World!",
+    "Goodbye, World!",
+    "Hello, Sky!",
+    "Goodbye, Volcano!",
+};
+}
+
 namespace extensions {
 namespace api {
 namespace cast_channel {
@@ -36,37 +48,52 @@
                                const MessageInfo& message));
 };
 
-class MockTCPClientSocket : public net::TCPClientSocket {
+class MockTCPSocket : public net::TCPClientSocket {
  public:
-  explicit MockTCPClientSocket() :
-      TCPClientSocket(net::AddressList(), NULL, net::NetLog::Source()) { }
-  virtual ~MockTCPClientSocket() { }
+  explicit MockTCPSocket(const net::MockConnect& connect_data) :
+      TCPClientSocket(net::AddressList(), NULL, net::NetLog::Source()),
+      connect_data_(connect_data) { }
 
-  MOCK_METHOD1(Connect, int(const net::CompletionCallback& callback));
-  MOCK_METHOD2(SetKeepAlive, bool(bool, int));
-  MOCK_METHOD1(SetNoDelay, bool(bool));
-  MOCK_METHOD3(Read, int(net::IOBuffer* buf, int buf_len,
-                         const net::CompletionCallback& callback));
-  MOCK_METHOD3(Write, int(net::IOBuffer* buf, int buf_len,
-                          const net::CompletionCallback& callback));
-  MOCK_METHOD0(Disconnect, void());
-};
+  virtual int Connect(const net::CompletionCallback& callback) OVERRIDE {
+    if (connect_data_.mode == net::ASYNC) {
+      CHECK_NE(connect_data_.result, net::ERR_IO_PENDING);
+      base::MessageLoop::current()->PostTask(
+          FROM_HERE,
+          base::Bind(callback, connect_data_.result));
+      return net::ERR_IO_PENDING;
+    } else {
+      return connect_data_.result;
+    }
+  }
 
-class MockSSLClientSocket : public net::MockClientSocket {
- public:
-  MockSSLClientSocket() : MockClientSocket(net::BoundNetLog()) { }
-  virtual ~MockSSLClientSocket() { }
+  virtual bool SetKeepAlive(bool a, int b) OVERRIDE {
+    // Always return true in tests

[Ryan Sleevi, 2013/12/12 00:48:14]

These comments provide zero value.

[Munjal (Google), 2013/12/12 17:53:47]

The reason I have these comments is so that the reader knows that the behavior
is intended. I find such comments useful in situations when the reader is not
sure if the person who wrote the code has made a mistake or is it intentional.

+    return true;
+  }
 
-  MOCK_METHOD1(Connect, int(const net::CompletionCallback& callback));
-  MOCK_METHOD3(Read, int(net::IOBuffer* buf, int buf_len,
-                         const net::CompletionCallback& callback));
-  MOCK_METHOD3(Write, int(net::IOBuffer* buf, int buf_len,
-                          const net::CompletionCallback& callback));
-  MOCK_METHOD0(Disconnect, void());
-  MOCK_CONST_METHOD0(WasEverUsed, bool());
-  MOCK_CONST_METHOD0(UsingTCPFastOpen, bool());
-  MOCK_CONST_METHOD0(WasNpnNegotiated, bool());
-  MOCK_METHOD1(GetSSLInfo, bool(net::SSLInfo*));
+  virtual bool SetNoDelay(bool b) OVERRIDE {
+    // Always return true in tests
+    return true;
+  }
+
+  virtual int Read(net::IOBuffer* buf, int buf_len,
+                   const net::CompletionCallback& callback) OVERRIDE {
+    // Always return IO_PENDING in tests
+    return net::ERR_IO_PENDING;
+  }
+
+  virtual int Write (net::IOBuffer* buf, int buf_len,
+                     const net::CompletionCallback& callback) OVERRIDE {
+    // Always return IO_PENDING in tests
+    return net::ERR_IO_PENDING;
+  }

[Ryan Sleevi, 2013/12/12 00:48:14]

This behaviour for Read and Write leaves me very uncomfortable with the actual
coverage quality of the unittests. Supporting the Mock Sockets is highly
desirable here so you can actually test your assumptinos about the underlying
sockets behaviour.

[Munjal (Google), 2013/12/12 17:53:47]

Actually the Read/Write methods of the TCP socket should never be called. I
changed the mock impl to reflect that.

+
+  virtual void Disconnect() OVERRIDE {
+    // Do nothing in tests
+  }
+
+ private:
+  net::MockConnect connect_data_;
 };
 
 class CompleteHandler {
@@ -95,72 +122,71 @@
 
   explicit TestCastSocket(MockCastSocketDelegate* delegate,
                           const std::string& url) :
-    CastSocket("abcdefg", GURL(url), delegate,
-               &capturing_net_log_),
-    mock_tcp_socket_(new MockTCPClientSocket()),
-    mock_ssl_socket_(new MockSSLClientSocket()),
-    owns_tcp_socket_(true),
-    owns_ssl_socket_(true),
-    extract_cert_result_(true),
-    send_auth_challenge_result_(net::ERR_IO_PENDING),
-    read_auth_challenge_reply_result_(net::ERR_IO_PENDING),
-    challenge_reply_result_(true) {
+      CastSocket("abcdefg", GURL(url), delegate,
+                 &capturing_net_log_),
+      ip_(CreateIPEndPoint()),
+      connect_index_(0),
+      extract_cert_result_(true),
+      challenge_reply_result_(true) {
   }
 
-  virtual ~TestCastSocket() {
-    if (owns_tcp_socket_) {
-      DCHECK(mock_tcp_socket_);
-      delete mock_tcp_socket_;
-    }
-    if (owns_ssl_socket_) {
-      DCHECK(mock_ssl_socket_);
-      delete mock_ssl_socket_;
-    }
+  static net::IPEndPoint CreateIPEndPoint() {
+    net::IPAddressNumber number;
+    number.push_back(192);
+    number.push_back(0);
+    number.push_back(0);
+    number.push_back(1);
+    return net::IPEndPoint(number, 8009);
   }
 
-  virtual void Close(const net::CompletionCallback& callback) OVERRIDE {
-    if (!owns_tcp_socket_)
-      mock_tcp_socket_ = NULL;
-    if (!owns_ssl_socket_)
-      mock_ssl_socket_ = NULL;
-    CastSocket::Close(callback);
+  virtual ~TestCastSocket() {
   }
 
-  void CreateNewSockets() {
-    owns_tcp_socket_ = true;
-    mock_tcp_socket_ = new MockTCPClientSocket();
-    owns_ssl_socket_ = true;
-    mock_ssl_socket_ = new MockSSLClientSocket();
+  // Helpers to set mock results for various operations.
+  void SetupTcp1Connect(net::IoMode mode, int result) {
+    tcp_connect_data_[0].reset(new net::MockConnect(mode, result));
   }
-
-  void SetExtractCertResult(bool value) {
-    extract_cert_result_ = value;
+  void SetupSsl1Connect(net::IoMode mode, int result) {
+    ssl_connect_data_[0].reset(new net::MockConnect(mode, result));
   }
-
-  void SetSendAuthChallengeResult(int result) {
-    send_auth_challenge_result_ = result;
+  void SetupTcp2Connect(net::IoMode mode, int result) {
+    tcp_connect_data_[1].reset(new net::MockConnect(mode, result));
   }
-
-  void SetReadAuthChallengeReplyResult(int result) {
-    read_auth_challenge_reply_result_ = result;
+  void SetupSsl2Connect(net::IoMode mode, int result) {
+    ssl_connect_data_[1].reset(new net::MockConnect(mode, result));
   }
-
+  void AddWriteResult(const net::MockWrite& write) {
+    writes_.push_back(write);
+  }
+  void AddReadResult(const net::MockRead& read) {
+    reads_.push_back(read);
+  }
+  void SetExtractCertResult(bool value) {
+    extract_cert_result_ = value;
+  }
   void SetChallengeReplyResult(bool value) {
     challenge_reply_result_ = value;
   }
 
-  MockTCPClientSocket* mock_tcp_socket_;
-  MockSSLClientSocket* mock_ssl_socket_;
-
  protected:
   virtual scoped_ptr<net::TCPClientSocket> CreateTcpSocket() OVERRIDE {
-    owns_tcp_socket_ = false;
-    return scoped_ptr<net::TCPClientSocket>(mock_tcp_socket_);
+    net::MockConnect* connect_data = tcp_connect_data_[connect_index_].get();
+    connect_data->peer_addr = ip_;
+    return scoped_ptr<net::TCPClientSocket>(new MockTCPSocket(*connect_data));
   }
 
-  virtual scoped_ptr<net::SSLClientSocket> CreateSslSocket() OVERRIDE {
-    owns_ssl_socket_ = false;
-    return scoped_ptr<net::SSLClientSocket>(mock_ssl_socket_);
+  virtual scoped_ptr<net::SSLClientSocket> CreateSslSocket(
+      scoped_ptr<net::StreamSocket> socket) OVERRIDE {
+    net::MockConnect* connect_data = ssl_connect_data_[connect_index_].get();
+    connect_data->peer_addr = ip_;
+    ++connect_index_;
+
+    ssl_data_.reset(new net::StaticSocketDataProvider(
+        reads_.data(), reads_.size(), writes_.data(), writes_.size()));
+    ssl_data_->set_connect_data(*connect_data);
+    return scoped_ptr<net::SSLClientSocket>(
+        new net::MockTCPClientSocket(
+            net::AddressList(), &capturing_net_log_, ssl_data_.get()));
   }
 
   virtual bool ExtractPeerCert(std::string* cert) OVERRIDE {
@@ -169,30 +195,24 @@
     return extract_cert_result_;
   }
 
-  virtual int SendAuthChallenge() OVERRIDE {
-    return send_auth_challenge_result_;
-  }
-
-  virtual int ReadAuthChallengeReply() OVERRIDE {
-    return read_auth_challenge_reply_result_;
-  }
-
   virtual bool VerifyChallengeReply() OVERRIDE {
     return challenge_reply_result_;
   }
 
  private:
   net::CapturingNetLog capturing_net_log_;
-  // Whether this object or the parent owns |mock_tcp_socket_|.
-  bool owns_tcp_socket_;
-  // Whether this object or the parent owns |mock_ssl_socket_|.
-  bool owns_ssl_socket_;
+  net::IPEndPoint ip_;
+  // Simulated connect data
+  scoped_ptr<net::MockConnect> tcp_connect_data_[2];
+  scoped_ptr<net::MockConnect> ssl_connect_data_[2];
+  // Simulated read / write data
+  std::vector<net::MockWrite> writes_;
+  std::vector<net::MockRead> reads_;
+  scoped_ptr<net::SocketDataProvider> ssl_data_;
+  // Number of times Connect method is called
+  size_t connect_index_;
   // Simulated result of peer cert extraction.
   bool extract_cert_result_;
-  // Simulated result to be returned by SendAuthChallenge.
-  int send_auth_challenge_result_;
-  // Simulated result to be returned by ReadAuthChallengeReply.
-  int read_auth_challenge_reply_result_;
   // Simulated result of verifying challenge reply.
   bool challenge_reply_result_;
 };
@@ -203,11 +223,32 @@
   virtual ~CastSocketTest() {}
 
   virtual void SetUp() OVERRIDE {
-    test_message_.namespace_ = "urn:test";
-    test_message_.source_id = "1";
-    test_message_.destination_id = "2";
-    test_message_.data.reset(new base::StringValue("Hello, World!"));
-    ASSERT_TRUE(MessageInfoToCastMessage(test_message_, &test_proto_));
+    // Create a few test messages
+    for (size_t i = 0; i < arraysize(test_messages_); i++) {
+      test_messages_[i].namespace_ = "urn:test";
+      test_messages_[i].source_id = "1";
+      test_messages_[i].destination_id = "2";
+      test_messages_[i].data.reset(new base::StringValue(kTestData[i]));
+      ASSERT_TRUE(MessageInfoToCastMessage(
+          test_messages_[i], &test_protos_[i]));
+      ASSERT_TRUE(CastSocket::Serialize(test_protos_[i], &test_proto_strs_[i]));
+    }
+
+    // Create a test auth request.
+    CastMessage request;
+    CreateAuthChallengeMessage(&request);
+    ASSERT_TRUE(CastSocket::Serialize(request, &auth_request_));
+
+    // Create a test auth reply.
+    MessageInfo reply;
+    reply.namespace_ = "urn:x-cast:com.google.cast.tp.deviceauth";
+    reply.source_id = "sender-0";
+    reply.destination_id = "receiver-0";
+    char D[] = "abcd";
+    reply.data.reset(base::BinaryValue::CreateWithCopiedBuffer(D, sizeof(D)));
+    CastMessage reply_msg;
+    ASSERT_TRUE(MessageInfoToCastMessage(reply, &reply_msg));
+    ASSERT_TRUE(CastSocket::Serialize(reply_msg, &auth_reply_));
   }
 
   virtual void TearDown() OVERRIDE {
@@ -224,85 +265,36 @@
     socket_ = TestCastSocket::CreateSecure(&mock_delegate_);
   }
 
-  // Sets an expectation that TCPClientSocket::Connect is called and
-  // returns |result| and stores the callback passed to it in |callback|.
-  void ExpectTcpConnect(net::CompletionCallback* callback, int result) {
-    EXPECT_CALL(mock_tcp_socket(), Connect(A<const net::CompletionCallback&>()))
-        .Times(1)
-        .WillOnce(DoAll(SaveArg<0>(callback), Return(result)));
-    EXPECT_CALL(mock_tcp_socket(), SetKeepAlive(_, _))
-        .WillOnce(Return(true));
-  }
-
-  // Same as ExpectTcpConnect but to return net::ERR_IO_PENDING.
-  void ExpectTcpConnectPending(net::CompletionCallback* callback) {
-    ExpectTcpConnect(callback, net::ERR_IO_PENDING);
-  }
-
-  // Sets an expectation that SSLClientSocket::Connect is called and
-  // returns |result| and stores the callback passed to it in |callback|.
-  void ExpectSslConnect(net::CompletionCallback* callback, int result) {
-    EXPECT_CALL(mock_ssl_socket(), Connect(A<const net::CompletionCallback&>()))
-        .Times(1)
-        .WillOnce(DoAll(SaveArg<0>(callback), Return(result)));
-  }
-
-  // Same as ExpectSslConnect but to return net::ERR_IO_PENDING.
-  void ExpectSslConnectPending(net::CompletionCallback* callback) {
-    ExpectSslConnect(callback, net::ERR_IO_PENDING);
-  }
-
-  // Sets an expectation that SSLClientSocket::Read is called |times| number
-  // of times and returns net::ERR_IO_PENDING.
-  void ExpectSslRead(int times) {
-    EXPECT_CALL(mock_ssl_socket(), Read(A<net::IOBuffer*>(),
-                                        A<int>(),
-                                        A<const net::CompletionCallback&>()))
-        .Times(times)
-        .WillOnce(Return(net::ERR_IO_PENDING));
-  }
-
   // Sets up CastSocket::Connect to succeed.
-  // Connecting the socket also starts the read loop; we expect the call to
-  // Read(), but never fire the read callback.
+  // Connecting the socket also starts the read loop; so we add a mock
+  // read result that returns IO_PENDING and callback is never fired.
   void ConnectHelper() {
-    net::CompletionCallback connect_callback1;
-    net::CompletionCallback connect_callback2;
+    socket_->SetupTcp1Connect(net::SYNCHRONOUS, net::OK);
+    socket_->SetupSsl1Connect(net::SYNCHRONOUS, net::OK);
+    socket_->AddReadResult(net::MockRead(net::ASYNC, net::ERR_IO_PENDING));
 
-    ExpectTcpConnect(&connect_callback1, net::OK);
-    ExpectSslConnect(&connect_callback2, net::OK);
     EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-    ExpectSslRead(1);
-
     socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
                                 base::Unretained(&handler_)));
-
-    EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
-    EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
+    RunPendingTasks();
   }
 
  protected:
-  MockTCPClientSocket& mock_tcp_socket() {
-    MockTCPClientSocket* mock_socket = socket_->mock_tcp_socket_;
-    DCHECK(mock_socket);
-    return *mock_socket;
+  // Runs all pending tasks in the message loop.
+  void RunPendingTasks() {
+    base::RunLoop run_loop;
+    run_loop.RunUntilIdle();
   }
 
-  MockSSLClientSocket& mock_ssl_socket() {
-    MockSSLClientSocket* mock_socket = socket_->mock_ssl_socket_;
-    DCHECK(mock_socket);
-    return *mock_socket;
-  }
-
-  void CallOnChallengeEvent(int result) {
-    socket_->OnChallengeEvent(result);
-  }
-
+  base::MessageLoop message_loop_;
   MockCastSocketDelegate mock_delegate_;
   scoped_ptr<TestCastSocket> socket_;
   CompleteHandler handler_;
-  MessageInfo test_message_;
-  CastMessage test_proto_;
+  MessageInfo test_messages_[arraysize(kTestData)];
+  CastMessage test_protos_[arraysize(kTestData)];
+  std::string test_proto_strs_[arraysize(kTestData)];
+  std::string auth_request_;
+  std::string auth_reply_;
 };
 
 // Tests URL parsing and validation.
@@ -335,6 +327,8 @@
 TEST_F(CastSocketTest, TestConnectAndClose) {
   CreateCastSocket();
   ConnectHelper();
+  EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
+  EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 
   EXPECT_CALL(handler_, OnCloseComplete(net::OK));
   socket_->Close(base::Bind(&CompleteHandler::OnCloseComplete,
@@ -348,19 +342,14 @@
 // - SSL connection succeeds (async)
 TEST_F(CastSocketTest, TestConnect) {
   CreateCastSocket();
+  socket_->SetupTcp1Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl1Connect(net::ASYNC, net::OK);
+  socket_->AddReadResult(net::MockRead(net::ASYNC, net::ERR_IO_PENDING));
 
-  net::CompletionCallback connect_callback1;
-  net::CompletionCallback connect_callback2;
-
-  ExpectTcpConnectPending(&connect_callback1);
-  ExpectSslConnectPending(&connect_callback2);
   EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-  ExpectSslRead(1);
-
   socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
                               base::Unretained(&handler_)));
-  connect_callback1.Run(net::OK);
-  connect_callback2.Run(net::OK);
+  RunPendingTasks();
 
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
@@ -374,35 +363,17 @@
 // - Second SSL connection succeeds (async)
 TEST_F(CastSocketTest, TestTwoStepConnect) {
   CreateCastSocket();
+  socket_->SetupTcp1Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl1Connect(net::ASYNC, net::ERR_CERT_AUTHORITY_INVALID);
+  socket_->SetupTcp2Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl2Connect(net::ASYNC, net::OK);
+  socket_->AddReadResult(net::MockRead(net::ASYNC, net::ERR_IO_PENDING));
 
-  // Expectations for the initial connect call
-  net::CompletionCallback tcp_connect_callback1;
-  net::CompletionCallback ssl_connect_callback1;
-
-  ExpectTcpConnectPending(&tcp_connect_callback1);
-  ExpectSslConnectPending(&ssl_connect_callback1);
-
-  // Start connect flow
+  EXPECT_CALL(handler_, OnConnectComplete(net::OK));
   socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
                               base::Unretained(&handler_)));
-  tcp_connect_callback1.Run(net::OK);
+  RunPendingTasks();
 
-  // Expectations for the second connect call
-  socket_->CreateNewSockets();
-  net::CompletionCallback tcp_connect_callback2;
-  net::CompletionCallback ssl_connect_callback2;
-  ExpectTcpConnectPending(&tcp_connect_callback2);
-  ExpectSslConnectPending(&ssl_connect_callback2);
-  EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-  ExpectSslRead(1);
-
-  // Trigger callbacks for the first connect
-  ssl_connect_callback1.Run(net::ERR_CERT_AUTHORITY_INVALID);
-
-  // Trigger callbacks for the second connect
-  tcp_connect_callback2.Run(net::OK);
-  ssl_connect_callback2.Run(net::OK);
-
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
@@ -416,35 +387,16 @@
 // - The flow should NOT be tried again
 TEST_F(CastSocketTest, TestMaxTwoConnectAttempts) {
   CreateCastSocket();
+  socket_->SetupTcp1Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl1Connect(net::ASYNC, net::ERR_CERT_AUTHORITY_INVALID);
+  socket_->SetupTcp2Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl2Connect(net::ASYNC, net::ERR_CERT_AUTHORITY_INVALID);
 
-  net::CompletionCallback tcp_connect_callback1;
-  net::CompletionCallback ssl_connect_callback1;
-
-  // Expectations for the initial connect call
-  ExpectTcpConnectPending(&tcp_connect_callback1);
-  ExpectSslConnectPending(&ssl_connect_callback1);
-
-  // Start connect flow
+  EXPECT_CALL(handler_, OnConnectComplete(net::ERR_CERT_AUTHORITY_INVALID));
   socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
                               base::Unretained(&handler_)));
-  tcp_connect_callback1.Run(net::OK);
+  RunPendingTasks();
 
-  socket_->CreateNewSockets();
-  net::CompletionCallback tcp_connect_callback2;
-  net::CompletionCallback ssl_connect_callback2;
-
-  // Expectations for the second connect call
-  ExpectTcpConnectPending(&tcp_connect_callback2);
-  ExpectSslConnectPending(&ssl_connect_callback2);
-  EXPECT_CALL(handler_, OnConnectComplete(net::ERR_CERT_AUTHORITY_INVALID));
-
-  // Trigger callbacks for the first connect
-  ssl_connect_callback1.Run(net::ERR_CERT_AUTHORITY_INVALID);
-
-  // Trigger callbacks for the second connect
-  tcp_connect_callback2.Run(net::OK);
-  ssl_connect_callback2.Run(net::ERR_CERT_AUTHORITY_INVALID);
-
   EXPECT_EQ(cast_channel::READY_STATE_CLOSED, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_CONNECT_ERROR, socket_->error_state());
 }
@@ -452,24 +404,16 @@
 // Test that when cert extraction fails the connection flow stops.
 TEST_F(CastSocketTest, TestCertExtractionFailure) {
   CreateCastSocket();
+  socket_->SetupTcp1Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl1Connect(net::ASYNC, net::ERR_CERT_AUTHORITY_INVALID);
+  // Set cert extraction to fail
+  socket_->SetExtractCertResult(false);
 
-  net::CompletionCallback connect_callback1;
-  net::CompletionCallback connect_callback2;
-
-  ExpectTcpConnectPending(&connect_callback1);
-  ExpectSslConnectPending(&connect_callback2);
-
+  EXPECT_CALL(handler_, OnConnectComplete(net::ERR_CERT_AUTHORITY_INVALID));
   socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
                               base::Unretained(&handler_)));
-  connect_callback1.Run(net::OK);
+  RunPendingTasks();
 
-  EXPECT_CALL(handler_, OnConnectComplete(net::ERR_CERT_AUTHORITY_INVALID));
-
-  // Set cert extraction to fail
-  socket_->SetExtractCertResult(false);
-  // Attempt to connect results in ERR_CERT_AUTHORTY_INVALID
-  connect_callback2.Run(net::ERR_CERT_AUTHORITY_INVALID);
-
   EXPECT_EQ(cast_channel::READY_STATE_CLOSED, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_CONNECT_ERROR, socket_->error_state());
 }
@@ -486,211 +430,186 @@
 TEST_F(CastSocketTest, TestFullSecureConnectionFlowAsync) {
   CreateCastSocketSecure();
 
-  net::CompletionCallback tcp_connect_callback1;
-  net::CompletionCallback ssl_connect_callback1;
+  socket_->SetupTcp1Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl1Connect(net::ASYNC, net::ERR_CERT_AUTHORITY_INVALID);
+  socket_->SetupTcp2Connect(net::ASYNC, net::OK);
+  socket_->SetupSsl2Connect(net::ASYNC, net::OK);
+  socket_->AddWriteResult(net::MockWrite(net::ASYNC, auth_request_.size()));
+  size_t body_size = auth_reply_.length() - 4;
+  const char* reply_data = auth_reply_.c_str();
+  socket_->AddReadResult(net::MockRead(net::ASYNC, reply_data, 4));
+  socket_->AddReadResult(net::MockRead(net::ASYNC, reply_data + 4, body_size));
+  socket_->AddReadResult(net::MockRead(net::ASYNC, net::ERR_IO_PENDING));
 
-  // Expectations for the initial connect call
-  ExpectTcpConnectPending(&tcp_connect_callback1);
-  ExpectSslConnectPending(&ssl_connect_callback1);
-
-  // Start connect flow
+  EXPECT_CALL(handler_, OnConnectComplete(net::OK));
   socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
                               base::Unretained(&handler_)));
-  tcp_connect_callback1.Run(net::OK);
+  RunPendingTasks();
 
-  socket_->CreateNewSockets();
-  net::CompletionCallback tcp_connect_callback2;
-  net::CompletionCallback ssl_connect_callback2;
-
-  // Expectations for the second connect call
-  ExpectTcpConnectPending(&tcp_connect_callback2);
-  ExpectSslConnectPending(&ssl_connect_callback2);
-  EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-
-  // Trigger callbacks for the first connect
-  ssl_connect_callback1.Run(net::ERR_CERT_AUTHORITY_INVALID);
-
-  // Trigger callbacks for the second connect
-  tcp_connect_callback2.Run(net::OK);
-  ssl_connect_callback2.Run(net::OK);
-
-  // Trigger callbacks for auth events.
-  CallOnChallengeEvent(net::OK);  // Sent challenge
-  CallOnChallengeEvent(net::OK);  // Received reply
-
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
 
-// Same as TestFullSecureConnectionFlowAsync, but operations are  synchronous.
+// Same as TestFullSecureConnectionFlowAsync, but operations are synchronous.
 TEST_F(CastSocketTest, TestFullSecureConnectionFlowSync) {
   CreateCastSocketSecure();
 
-  net::CompletionCallback tcp_connect_callback;
-  net::CompletionCallback ssl_connect_callback;
+  socket_->SetupTcp1Connect(net::SYNCHRONOUS, net::OK);
+  socket_->SetupSsl1Connect(net::SYNCHRONOUS, net::ERR_CERT_AUTHORITY_INVALID);
+  socket_->SetupTcp2Connect(net::SYNCHRONOUS, net::OK);
+  socket_->SetupSsl2Connect(net::SYNCHRONOUS, net::OK);
+  socket_->AddWriteResult(net::MockWrite(
+      net::SYNCHRONOUS, auth_request_.size()));
+  size_t body_size = auth_reply_.length() - 4;
+  const char* reply_data = auth_reply_.c_str();
+  socket_->AddReadResult(net::MockRead(net::SYNCHRONOUS, reply_data, 4));
+  socket_->AddReadResult(net::MockRead(
+      net::SYNCHRONOUS, reply_data + 4, body_size));
+  socket_->AddReadResult(net::MockRead(net::ASYNC, net::ERR_IO_PENDING));
 
-  // Expectations for the connect calls
-  ExpectTcpConnect(&tcp_connect_callback, net::OK);
-  ExpectSslConnect(&ssl_connect_callback, net::OK);
   EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-
-  socket_->SetSendAuthChallengeResult(net::OK);
-  socket_->SetReadAuthChallengeReplyResult(net::OK);
-
-  // Start connect flow
   socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
                               base::Unretained(&handler_)));
+  RunPendingTasks();
 
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
 
-// Tests writing a single message where the completion is signaled via
-// callback.
-TEST_F(CastSocketTest, TestWriteViaCallback) {
+// Tests writing a single message - asynchronous.
+TEST_F(CastSocketTest, TestWriteAsync) {
   CreateCastSocket();
+
+  socket_->AddWriteResult(net::MockWrite(net::ASYNC, 39));
   ConnectHelper();
 
-  net::CompletionCallback write_callback;
-
-  EXPECT_CALL(mock_ssl_socket(),
-              Write(A<net::IOBuffer*>(),
-                    39,
-                    A<const net::CompletionCallback&>()))
-    .Times(1)
-    .WillOnce(DoAll(SaveArg<2>(&write_callback),
-                    Return(net::ERR_IO_PENDING)));
   EXPECT_CALL(handler_, OnWriteComplete(39));
-  socket_->SendMessage(test_message_,
+  socket_->SendMessage(test_messages_[0],
                        base::Bind(&CompleteHandler::OnWriteComplete,
                                   base::Unretained(&handler_)));
-  write_callback.Run(39);
+  RunPendingTasks();
+
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
 
-// Tests writing a single message where the Write() returns directly.
-TEST_F(CastSocketTest, TestWrite) {
+// Tests writing a single message - synchronous.
+TEST_F(CastSocketTest, TestWriteSync) {
   CreateCastSocket();
+
+  const std::string& msg = test_proto_strs_[0];
+  size_t msg_size = msg.size();
+  socket_->AddWriteResult(net::MockWrite(net::SYNCHRONOUS, msg_size));
   ConnectHelper();
 
-  EXPECT_CALL(mock_ssl_socket(),
-              Write(A<net::IOBuffer*>(),
-                    39,
-                    A<const net::CompletionCallback&>()))
-    .Times(1)
-    .WillOnce(Return(39));
-  EXPECT_CALL(handler_, OnWriteComplete(39));
-  socket_->SendMessage(test_message_,
-                      base::Bind(&CompleteHandler::OnWriteComplete,
-                                 base::Unretained(&handler_)));
+  EXPECT_CALL(handler_, OnWriteComplete(msg_size));
+  socket_->SendMessage(test_messages_[0],
+                       base::Bind(&CompleteHandler::OnWriteComplete,
+                                  base::Unretained(&handler_)));
+  RunPendingTasks();
+
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
 
-// Tests writing multiple messages.
+// Tests writing multiple messages - asynchronous.
 TEST_F(CastSocketTest, TestWriteMany) {
   CreateCastSocket();
-  ConnectHelper();
-  std::string messages[4];
-  messages[0] = "Hello, World!";
-  messages[1] = "Goodbye, World!";
-  messages[2] = "Hello, Sky!";
-  messages[3] = "Goodbye, Volcano!";
-  int sizes[4] = {39, 41, 37, 43};
-  MessageInfo message_info[4];
-  net::CompletionCallback write_callback;
 
-  for (int i = 0; i < 4; i++) {
-    EXPECT_CALL(mock_ssl_socket(),
-                Write(A<net::IOBuffer*>(),
-                      sizes[i],
-                      A<const net::CompletionCallback&>()))
-      .WillRepeatedly(DoAll(SaveArg<2>(&write_callback),
-                            Return(net::ERR_IO_PENDING)));
-    EXPECT_CALL(handler_, OnWriteComplete(sizes[i]));
+  for (size_t i = 0; i < arraysize(test_messages_); i++) {
+    size_t msg_size = test_proto_strs_[i].size();
+    socket_->AddWriteResult(net::MockWrite(net::ASYNC, msg_size));
+    EXPECT_CALL(handler_, OnWriteComplete(msg_size));
   }
+  ConnectHelper();
 
-  for (int i = 0; i < 4; i++) {
-    message_info[i].namespace_ = "urn:test";
-    message_info[i].source_id = "1";
-    message_info[i].destination_id = "2";
-    message_info[i].data.reset(new base::StringValue(messages[i]));
-    socket_->SendMessage(message_info[i],
+  for (size_t i = 0; i < arraysize(test_messages_); i++) {
+    socket_->SendMessage(test_messages_[i],
                          base::Bind(&CompleteHandler::OnWriteComplete,
                                     base::Unretained(&handler_)));
   }
-  for (int i = 0; i < 4; i++) {
-    write_callback.Run(sizes[i]);
-  }
+  RunPendingTasks();
+
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
 
-// Tests error on writing.
-TEST_F(CastSocketTest, TestWriteError) {
+// Tests error on writing - synchronous.
+TEST_F(CastSocketTest, TestWriteErrorSync) {
   CreateCastSocket();
+  socket_->AddWriteResult(net::MockWrite(net::SYNCHRONOUS,
+                                         net::ERR_SOCKET_NOT_CONNECTED));
   ConnectHelper();
-  net::CompletionCallback write_callback;
 
-  EXPECT_CALL(mock_ssl_socket(),
-              Write(A<net::IOBuffer*>(),
-                    39,
-                    A<const net::CompletionCallback&>()))
-    .Times(1)
-    .WillOnce(DoAll(SaveArg<2>(&write_callback),
-                    Return(net::ERR_SOCKET_NOT_CONNECTED)));
   EXPECT_CALL(handler_, OnWriteComplete(net::ERR_SOCKET_NOT_CONNECTED));
   EXPECT_CALL(mock_delegate_,
               OnError(socket_.get(), cast_channel::CHANNEL_ERROR_SOCKET_ERROR));
-  socket_->SendMessage(test_message_,
+  socket_->SendMessage(test_messages_[0],
                        base::Bind(&CompleteHandler::OnWriteComplete,
                                   base::Unretained(&handler_)));
+  RunPendingTasks();
+
   EXPECT_EQ(cast_channel::READY_STATE_CLOSED, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_SOCKET_ERROR, socket_->error_state());
 }
 
-// Tests reading a single message.
-TEST_F(CastSocketTest, TestRead) {
+// Tests error on writing - asynchronous.
+TEST_F(CastSocketTest, TestWriteErrorAsync) {
   CreateCastSocket();
+  socket_->AddWriteResult(net::MockWrite(net::ASYNC,
+                                         net::ERR_SOCKET_NOT_CONNECTED));
+  ConnectHelper();
 
-  net::CompletionCallback connect_callback1;
-  net::CompletionCallback connect_callback2;
-  net::CompletionCallback read_callback;
+  EXPECT_CALL(handler_, OnWriteComplete(net::ERR_SOCKET_NOT_CONNECTED));
+  EXPECT_CALL(mock_delegate_,
+              OnError(socket_.get(), cast_channel::CHANNEL_ERROR_SOCKET_ERROR));
+  socket_->SendMessage(test_messages_[0],
+                       base::Bind(&CompleteHandler::OnWriteComplete,
+                                  base::Unretained(&handler_)));
+  RunPendingTasks();
 
-  std::string message_data;
-  ASSERT_TRUE(CastSocket::Serialize(test_proto_, &message_data));
+  EXPECT_EQ(cast_channel::READY_STATE_CLOSED, socket_->ready_state());
+  EXPECT_EQ(cast_channel::CHANNEL_ERROR_SOCKET_ERROR, socket_->error_state());
+}
 
-  ExpectTcpConnect(&connect_callback1, net::ERR_IO_PENDING);
-  ExpectSslConnect(&connect_callback2, net::ERR_IO_PENDING);
+// Test that when an error occurrs in one write, write callback is invoked for
+// all pending writes with the error.
+TEST_F(CastSocketTest, TestWriteErrorWithMultiplePendingWritesAsync) {
+  CreateCastSocket();
+  socket_->AddWriteResult(net::MockWrite(net::ASYNC,
+                                         net::ERR_SOCKET_NOT_CONNECTED));
+  ConnectHelper();
 
-  EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-  EXPECT_CALL(mock_ssl_socket(), Read(A<net::IOBuffer*>(),
-                                      A<int>(),
-                                      A<const net::CompletionCallback&>()))
-      .Times(3)
-      .WillRepeatedly(DoAll(SaveArg<2>(&read_callback),
-                            Return(net::ERR_IO_PENDING)));
+  const int num_writes = arraysize(test_messages_);
+  EXPECT_CALL(handler_, OnWriteComplete(net::ERR_SOCKET_NOT_CONNECTED))
+      .Times(num_writes);
+  EXPECT_CALL(mock_delegate_,
+              OnError(socket_.get(), cast_channel::CHANNEL_ERROR_SOCKET_ERROR));
+  for (int i = 0; i < num_writes; i++) {
+    socket_->SendMessage(test_messages_[i],
+                         base::Bind(&CompleteHandler::OnWriteComplete,
+                                    base::Unretained(&handler_)));
+  }
+  RunPendingTasks();
 
-  // Expect the test message to be read and invoke the delegate.
+  EXPECT_EQ(cast_channel::READY_STATE_CLOSED, socket_->ready_state());
+  EXPECT_EQ(cast_channel::CHANNEL_ERROR_SOCKET_ERROR, socket_->error_state());
+}
+
+// Tests reading a single message.
+TEST_F(CastSocketTest, TestRead) {
+  CreateCastSocket();
+
+  const std::string& message = test_proto_strs_[0];
+  const char* data = message.c_str();
+  size_t body_size = message.size() - 4;
+  socket_->AddReadResult(net::MockRead(net::ASYNC, data, 4));
+  socket_->AddReadResult(net::MockRead(net::ASYNC, data + 4, body_size));
   EXPECT_CALL(mock_delegate_,
               OnMessage(socket_.get(), A<const MessageInfo&>()));
 
-  // Connect the socket.
-  socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
-                              base::Unretained(&handler_)));
-  connect_callback1.Run(net::OK);
-  connect_callback2.Run(net::OK);
+  ConnectHelper();
 
-  // Put the test header and message into the io_buffers and invoke the read
-  // callbacks.
-  memcpy(socket_->header_read_buffer_->StartOfBuffer(),
-         message_data.c_str(), 4);
-  read_callback.Run(4);
-  memcpy(socket_->body_read_buffer_->StartOfBuffer(),
-         message_data.c_str() + 4, 35);
-  read_callback.Run(35);
-
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
@@ -699,57 +618,20 @@
 TEST_F(CastSocketTest, TestReadMany) {
   CreateCastSocket();
 
-  net::CompletionCallback connect_callback1;
-  net::CompletionCallback connect_callback2;
-  net::CompletionCallback read_callback;
-
-  std::string messages[4];
-  messages[0] = "Hello, World!";
-  messages[1] = "Goodbye, World!";
-  messages[2] = "Hello, Sky!";
-  messages[3] = "Goodbye, Volcano!";
-  int sizes[4] = {35, 37, 33, 39};
-  std::string message_data[4];
-
-  // Set up test data
-  for (int i = 0; i < 4; i++) {
-    test_proto_.set_payload_utf8(messages[i]);
-    ASSERT_TRUE(CastSocket::Serialize(test_proto_, &message_data[i]));
+  size_t num_reads = arraysize(test_proto_strs_);
+  for (size_t i = 0; i < num_reads; i++) {
+    const char* data = test_proto_strs_[i].c_str();
+    size_t body_size = test_proto_strs_[i].size() - 4;
+    socket_->AddReadResult(net::MockRead(net::ASYNC, data, 4));
+    socket_->AddReadResult(net::MockRead(net::ASYNC, data + 4, body_size));
   }
 
-  ExpectTcpConnect(&connect_callback1, net::ERR_IO_PENDING);
-  ExpectSslConnect(&connect_callback2, net::ERR_IO_PENDING);
+  EXPECT_CALL(mock_delegate_,
+              OnMessage(socket_.get(), A<const MessageInfo&>()))
+      .Times(num_reads);
 
-  EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-  EXPECT_CALL(mock_ssl_socket(), Read(A<net::IOBuffer*>(),
-                                      A<int>(),
-                                      A<const net::CompletionCallback&>()))
-    .Times(9)
-    .WillRepeatedly(DoAll(SaveArg<2>(&read_callback),
-                          Return(net::ERR_IO_PENDING)));
+  ConnectHelper();
 
-  // Expect the test messages to be read and invoke the delegate.
-  EXPECT_CALL(mock_delegate_, OnMessage(socket_.get(),
-                                        A<const MessageInfo&>()))
-    .Times(4);
-
-  // Connect the socket.
-  socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
-                              base::Unretained(&handler_)));
-  connect_callback1.Run(net::OK);
-  connect_callback2.Run(net::OK);
-
-  // Put the test headers and messages into the io_buffer and invoke the read
-  // callbacks.
-  for (int i = 0; i < 4; i++) {
-    memcpy(socket_->header_read_buffer_->StartOfBuffer(),
-           message_data[i].c_str(), 4);
-    read_callback.Run(4);
-    memcpy(socket_->body_read_buffer_->StartOfBuffer(),
-           message_data[i].c_str() + 4, sizes[i]);
-    read_callback.Run(sizes[i]);
-  }
-
   EXPECT_EQ(cast_channel::READY_STATE_OPEN, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_NONE, socket_->error_state());
 }
@@ -757,32 +639,13 @@
 // Tests error on reading.
 TEST_F(CastSocketTest, TestReadError) {
   CreateCastSocket();
+  socket_->AddReadResult(net::MockRead(
+      net::ASYNC, net::ERR_SOCKET_NOT_CONNECTED));
 
-  net::CompletionCallback connect_callback1;
-  net::CompletionCallback connect_callback2;
-  net::CompletionCallback read_callback;
-
-  ExpectTcpConnect(&connect_callback1, net::ERR_IO_PENDING);
-  ExpectSslConnect(&connect_callback2, net::ERR_IO_PENDING);
-
-  EXPECT_CALL(handler_, OnConnectComplete(net::OK));
-  EXPECT_CALL(mock_ssl_socket(), Read(A<net::IOBuffer*>(),
-                                      A<int>(),
-                                      A<const net::CompletionCallback&>()))
-    .WillOnce(DoAll(SaveArg<2>(&read_callback),
-                    Return(net::ERR_IO_PENDING)));
   EXPECT_CALL(mock_delegate_,
               OnError(socket_.get(), cast_channel::CHANNEL_ERROR_SOCKET_ERROR));
+  ConnectHelper();
 
-  // Connect the socket.
-  socket_->Connect(base::Bind(&CompleteHandler::OnConnectComplete,
-                              base::Unretained(&handler_)));
-  connect_callback1.Run(net::OK);
-  connect_callback2.Run(net::OK);
-
-  // Cause an error.
-  read_callback.Run(net::ERR_SOCKET_NOT_CONNECTED);
-
   EXPECT_EQ(cast_channel::READY_STATE_CLOSED, socket_->ready_state());
   EXPECT_EQ(cast_channel::CHANNEL_ERROR_SOCKET_ERROR, socket_->error_state());
 }