3 files changed, 95 insertions, 18 deletions

diff --git a/net/socket/socket_test_util.cc b/net/socket/socket_test_util.cc
index 577c7d2..ab10549 100644
--- a/net/socket/socket_test_util.cc
+++ b/net/socket/socket_test_util.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -786,6 +786,12 @@ bool MockTCPClientSocket::IsConnectedAndIdle() const {
   return IsConnected();
 }
 
+int MockTCPClientSocket::GetPeerAddress(AddressList* address) const {
+  if (!IsConnected())
+    return ERR_SOCKET_NOT_CONNECTED;
+  return MockClientSocket::GetPeerAddress(address);
+}
+
 bool MockTCPClientSocket::WasEverUsed() const {
   return was_used_to_convey_data_;
 }
diff --git a/net/socket/socket_test_util.h b/net/socket/socket_test_util.h
index 823c3b8..42d4a46 100644
--- a/net/socket/socket_test_util.h
+++ b/net/socket/socket_test_util.h
@@ -622,6 +622,7 @@ class MockTCPClientSocket : public MockClientSocket {
   virtual void Disconnect();
   virtual bool IsConnected() const;
   virtual bool IsConnectedAndIdle() const;
+  virtual int GetPeerAddress(AddressList* address) const;
   virtual bool WasEverUsed() const;
   virtual bool UsingTCPFastOpen() const;
 
diff --git a/net/socket/transport_client_socket_unittest.cc b/net/socket/transport_client_socket_unittest.cc
index 837749e..28b70d4 100644
--- a/net/socket/transport_client_socket_unittest.cc
+++ b/net/socket/transport_client_socket_unittest.cc
@@ -48,9 +48,7 @@ class TransportClientSocketTest
   }
   virtual void DidRead(ListenSocket*, const char* str, int len) {
     // TODO(dkegel): this might not be long enough to tickle some bugs.
-    connected_sock_->Send(kServerReply,
-                          arraysize(kServerReply) - 1,
-                          false /* don't append line feed */);
+    SendServerReply();
   }
   virtual void DidClose(ListenSocket* sock) {}
 
@@ -70,6 +68,16 @@ class TransportClientSocketTest
     connected_sock_->ResumeReads();
   }
 
+  void SendServerReply() {
+    connected_sock_->Send(kServerReply, arraysize(kServerReply) - 1,
+                          false /* Don't append line feed */);
+  }
+
+  uint32 DrainClientSocket(IOBuffer* buf,
+                           uint32 buf_len,
+                           uint32 bytes_to_read,
+                           TestCompletionCallback* callback);
+
  protected:
   int listen_port_;
   CapturingNetLog net_log_;
@@ -115,6 +123,26 @@ void TransportClientSocketTest::SetUp() {
                                                    NetLog::Source()));
 }
 
+uint32 TransportClientSocketTest::DrainClientSocket(
+    IOBuffer* buf, uint32 buf_len,
+    uint32 bytes_to_read, TestCompletionCallback* callback) {
+  int rv = OK;
+  uint32 bytes_read = 0;
+
+  while (bytes_read < bytes_to_read) {
+    rv = sock_->Read(buf, buf_len, callback);
+    EXPECT_TRUE(rv >= 0 || rv == ERR_IO_PENDING);
+
+    if (rv == ERR_IO_PENDING)
+      rv = callback->WaitForResult();
+
+    EXPECT_GE(rv, 0);
+    bytes_read += rv;
+  }
+
+  return bytes_read;
+}
+
 // TODO(leighton):  Add SCTP to this list when it is ready.
 INSTANTIATE_TEST_CASE_P(ClientSocket,
                         TransportClientSocketTest,
@@ -147,9 +175,59 @@ TEST_P(TransportClientSocketTest, Connect) {
   EXPECT_FALSE(sock_->IsConnected());
 }
 
-// TODO(wtc): Add unit tests for IsConnectedAndIdle:
-//   - Server closes a connection.
-//   - Server sends data unexpectedly.
+TEST_P(TransportClientSocketTest, IsConnected) {
+  scoped_refptr<IOBuffer> buf(new IOBuffer(4096));
+  TestCompletionCallback callback;
+  uint32 bytes_read;
+
+  EXPECT_FALSE(sock_->IsConnected());
+  int rv = sock_->Connect(&callback);
+  if (rv != OK) {
+    ASSERT_EQ(rv, ERR_IO_PENDING);
+
+    rv = callback.WaitForResult();
+    EXPECT_EQ(rv, OK);
+  }
+
+  EXPECT_TRUE(sock_->IsConnected());
+  EXPECT_TRUE(sock_->IsConnectedAndIdle());
+
+  // Server sends some data
+  SendServerReply();
+  // Drain a single byte so we know we've received some data.
+  bytes_read = DrainClientSocket(buf, 1, 1, &callback);
+  ASSERT_EQ(bytes_read, 1u);
+
+  // Socket should be considered connected, but not idle,
+  EXPECT_TRUE(sock_->IsConnected());
+  EXPECT_FALSE(sock_->IsConnectedAndIdle());
+
+  bytes_read = DrainClientSocket(buf, 4096, arraysize(kServerReply) - 2,
+                                 &callback);
+  ASSERT_EQ(bytes_read, arraysize(kServerReply) - 2);
+
+  // After draining the data, the socket should be back to connected
+  // and idle.
+  EXPECT_TRUE(sock_->IsConnected());
+  EXPECT_TRUE(sock_->IsConnectedAndIdle());
+
+  // This time close the server socket immediately after sending.
+  SendServerReply();
+  CloseServerSocket();
+
+  // As above because of data.
+  EXPECT_TRUE(sock_->IsConnected());
+  EXPECT_FALSE(sock_->IsConnectedAndIdle());
+
+  bytes_read = DrainClientSocket(buf, 4096, arraysize(kServerReply) - 1,
+                                 &callback);
+  ASSERT_EQ(bytes_read, arraysize(kServerReply) - 1);
+
+  // Once the data is drained, the socket should now be seen as
+  // closed.
+  EXPECT_FALSE(sock_->IsConnected());
+  EXPECT_FALSE(sock_->IsConnectedAndIdle());
+}
 
 TEST_P(TransportClientSocketTest, Read) {
   TestCompletionCallback callback;
@@ -175,17 +253,9 @@ TEST_P(TransportClientSocketTest, Read) {
   }
 
   scoped_refptr<IOBuffer> buf(new IOBuffer(4096));
-  uint32 bytes_read = 0;
-  while (bytes_read < arraysize(kServerReply) - 1) {
-    rv = sock_->Read(buf, 4096, &callback);
-    EXPECT_TRUE(rv >= 0 || rv == ERR_IO_PENDING);
-
-    if (rv == ERR_IO_PENDING)
-      rv = callback.WaitForResult();
-
-    ASSERT_GE(rv, 0);
-    bytes_read += rv;
-  }
+  uint32 bytes_read = DrainClientSocket(buf, 4096, arraysize(kServerReply) - 1,
+                                        &callback);
+  ASSERT_EQ(bytes_read, arraysize(kServerReply) - 1);
 
   // All data has been read now.  Read once more to force an ERR_IO_PENDING, and
   // then close the server socket, and note the close.