1 files changed, 18 insertions, 3 deletions
diff --git a/net/socket/socket_test_util.cc b/net/socket/socket_test_util.cc
index b1af59f..4b6a3ad 100644
--- a/net/socket/socket_test_util.cc
+++ b/net/socket/socket_test_util.cc
@@ -290,17 +290,32 @@ void MockClientSocketFactory::ResetNextMockIndexes() {
   mock_ssl_sockets_.ResetNextIndex();
 }
 
+ClientSocket* MockClientSocketFactory::GetMockTCPClientSocket(int index) const {
+  return tcp_client_sockets_[index];
+}
+
+SSLClientSocket* MockClientSocketFactory::GetMockSSLClientSocket(
+    int index) const {
+  return ssl_client_sockets_[index];
+}
+
 ClientSocket* MockClientSocketFactory::CreateTCPClientSocket(
     const AddressList& addresses) {
-  return new MockTCPClientSocket(addresses, mock_sockets_.GetNext());
+  ClientSocket* socket =
+      new MockTCPClientSocket(addresses, mock_sockets_.GetNext());
+  tcp_client_sockets_.push_back(socket);
+  return socket;
 }
 
 SSLClientSocket* MockClientSocketFactory::CreateSSLClientSocket(
     ClientSocket* transport_socket,
     const std::string& hostname,
     const SSLConfig& ssl_config) {
-  return new MockSSLClientSocket(transport_socket, hostname, ssl_config,
-                                 mock_ssl_sockets_.GetNext());
+  SSLClientSocket* socket =
+      new MockSSLClientSocket(transport_socket, hostname, ssl_config,
+                              mock_ssl_sockets_.GetNext());
+  ssl_client_sockets_.push_back(socket);
+  return socket;
 }
 
 int TestSocketRequest::WaitForResult() {