Keep pool of pre-connected DNS sockets

net/dns/dns_session_unittest.cc

Index: net/dns/dns_session_unittest.cc
diff --git a/net/dns/dns_session_unittest.cc b/net/dns/dns_session_unittest.cc
new file mode 100644
index 0000000000000000000000000000000000000000..28d9717e1631311d6327e3300546f523e588a9ff
--- /dev/null
+++ b/net/dns/dns_session_unittest.cc
@@ -0,0 +1,223 @@
+// Copyright (c) 2012 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/dns/dns_session.h"
+
+#include <list>
+
+#include "base/bind.h"
+#include "base/memory/scoped_ptr.h"
+#include "base/rand_util.h"
+#include "base/stl_util.h"
+#include "net/base/net_log.h"
+#include "net/dns/dns_protocol.h"
+#include "net/dns/dns_socket_pool.h"
+#include "net/socket/socket_test_util.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace net {
+
+namespace {
+
+class TestClientSocketFactory : public ClientSocketFactory {
+ public:
+  virtual ~TestClientSocketFactory();
+
+  virtual DatagramClientSocket* CreateDatagramClientSocket(
+      DatagramSocket::BindType bind_type,
+      const RandIntCallback& rand_int_cb,
+      net::NetLog* net_log,
+      const net::NetLog::Source& source) OVERRIDE;
+
+  virtual StreamSocket* CreateTransportClientSocket(
+      const AddressList& addresses,
+      NetLog*, const NetLog::Source&) {

[szym, 2012/09/26 11:16:12]

nit: OVERRIDE

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

Done.

+    NOTIMPLEMENTED();
+    return NULL;
+  }
+
+  virtual SSLClientSocket* CreateSSLClientSocket(
+      ClientSocketHandle* transport_socket,
+      const HostPortPair& host_and_port,
+      const SSLConfig& ssl_config,
+      const SSLClientSocketContext& context) {

[szym, 2012/09/26 11:16:12]

nit: OVERRIDE

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

Done.

+    NOTIMPLEMENTED();
+    return NULL;
+  }
+
+  virtual void ClearSSLSessionCache() { NOTIMPLEMENTED(); }

[szym, 2012/09/26 11:16:12]

nit: OVERRIDE

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

Done.

+
+ private:
+  std::list<SocketDataProvider*> data_providers_;
+};
+
+struct PoolEvent {
+  enum { ALLOCATE, FREE } action;
+  unsigned server_index;
+};
+
+class DnsSessionTest : public testing::Test {
+ public:
+  void OnSocketAllocated(unsigned server_index);
+  void OnSocketFreed(unsigned server_index);
+
+ protected:
+  void Initialize(unsigned num_servers);
+  scoped_ptr<DnsSession::SocketLease> Allocate(unsigned server_index);
+  bool DidAllocate(unsigned server_index);
+  bool DidFree(unsigned server_index);
+  bool NoMoreEvents();
+
+  DnsConfig config_;
+  scoped_ptr<TestClientSocketFactory> test_client_socket_factory_;
+  scoped_refptr<DnsSession> session_;
+  NetLog::Source source_;
+
+ private:
+  bool ExpectEvent(const PoolEvent& event);
+  std::list<PoolEvent> events_;
+};
+
+class NullDnsSocketPool : public DnsSocketPool {

[szym, 2012/09/26 11:16:12]

I suggest Mock rather than Null.

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

Done.

+ public:
+  NullDnsSocketPool(ClientSocketFactory* factory, DnsSessionTest* test)
+     : DnsSocketPool(factory), test_(test) { }
+
+  virtual ~NullDnsSocketPool() { }
+
+  virtual void Initialize(
+      NetLog* net_log,
+      const std::vector<IPEndPoint>* nameservers) {
+    InitializeInternal(net_log, nameservers);
+  }
+
+  virtual scoped_ptr<DatagramClientSocket> AllocateSocket(
+      unsigned server_index) OVERRIDE {
+    test_->OnSocketAllocated(server_index);
+    return CreateConnectedSocket(server_index).Pass();

[szym, 2012/09/26 11:16:12]

No need for Pass().

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

Done.

+  }
+
+  virtual void FreeSocket(
+      unsigned server_index,
+      scoped_ptr<DatagramClientSocket> socket) OVERRIDE {
+    test_->OnSocketFreed(server_index);
+  }
+
+ private:
+  DnsSessionTest* test_;
+};
+
+void DnsSessionTest::Initialize(unsigned num_servers) {
+  CHECK(num_servers < 256u);
+  config_.nameservers.clear();
+  IPAddressNumber dns_ip;
+  bool rv = ParseIPLiteralToNumber("192.168.1.0", &dns_ip);
+  EXPECT_TRUE(rv);
+  for (unsigned char i = 0; i < num_servers; ++i) {
+    dns_ip[3] = i;
+    IPEndPoint dns_endpoint(dns_ip, dns_protocol::kDefaultPort);
+    config_.nameservers.push_back(dns_endpoint);
+  }
+
+  test_client_socket_factory_.reset(new TestClientSocketFactory());
+
+  DnsSocketPool* dns_socket_pool =
+      new NullDnsSocketPool(test_client_socket_factory_.get(), this);
+
+  session_ = new DnsSession(config_,
+                            scoped_ptr<DnsSocketPool>(dns_socket_pool),
+                            base::Bind(&base::RandInt),
+                            NULL /* NetLog */);
+
+  events_.clear();
+}
+
+scoped_ptr<DnsSession::SocketLease> DnsSessionTest::Allocate(
+    unsigned server_index) {
+  return session_->AllocateSocket(server_index, source_);
+}
+
+bool DnsSessionTest::DidAllocate(unsigned server_index) {
+  PoolEvent expected_event = { PoolEvent::ALLOCATE, server_index };
+  return ExpectEvent(expected_event);
+}
+
+bool DnsSessionTest::DidFree(unsigned server_index) {
+  PoolEvent expected_event = { PoolEvent::FREE, server_index };
+  return ExpectEvent(expected_event);
+}
+
+bool DnsSessionTest::NoMoreEvents() {
+  return events_.empty();
+}
+
+void DnsSessionTest::OnSocketAllocated(unsigned server_index) {
+  PoolEvent event = { PoolEvent::ALLOCATE, server_index };
+  events_.push_back(event);
+}
+
+void DnsSessionTest::OnSocketFreed(unsigned server_index) {
+  PoolEvent event = { PoolEvent::FREE, server_index };
+  events_.push_back(event);
+}
+
+bool DnsSessionTest::ExpectEvent(const PoolEvent& expected) {
+  if (events_.empty()) {
+    return false;
+  }
+
+  const PoolEvent actual = events_.front();
+  if ((expected.action != actual.action)
+      || (expected.server_index != actual.server_index)) {
+    return false;
+  }
+  events_.pop_front();
+
+  return true;
+}
+
+DatagramClientSocket* TestClientSocketFactory::CreateDatagramClientSocket(
+    DatagramSocket::BindType bind_type,
+    const RandIntCallback& rand_int_cb,
+    net::NetLog* net_log,
+    const net::NetLog::Source& source) {
+  // We're not actually expecting to send or receive any data, so use the
+  // simplest SocketDataProvider with no data supplied.
+  SocketDataProvider* data_provider = new StaticSocketDataProvider();
+  data_providers_.push_back(data_provider);
+  MockUDPClientSocket* socket = new MockUDPClientSocket(data_provider, net_log);
+  data_provider->set_socket(socket);
+  return socket;
+}
+
+TestClientSocketFactory::~TestClientSocketFactory() {
+  STLDeleteElements(&data_providers_);
+}
+
+TEST_F(DnsSessionTest, AllocateFree) {
+  scoped_ptr<DnsSession::SocketLease> lease1, lease2;
+
+  Initialize(2);
+  EXPECT_TRUE(NoMoreEvents());
+
+  lease1 = Allocate(0);
+  EXPECT_TRUE(DidAllocate(0));
+  EXPECT_TRUE(NoMoreEvents());
+
+  lease2 = Allocate(1);
+  EXPECT_TRUE(DidAllocate(1));
+  EXPECT_TRUE(NoMoreEvents());
+
+  lease1.reset();
+  EXPECT_TRUE(DidFree(0));
+  EXPECT_TRUE(NoMoreEvents());
+
+  lease2.reset();
+  EXPECT_TRUE(DidFree(1));
+  EXPECT_TRUE(NoMoreEvents());

[szym, 2012/09/26 11:16:12]

Since all that's tested here is whether DnsSession and SocketLease transparently
execute Allocate/Free calls on the DnsSocketPool, I wonder if using gmock
wouldn't be more suitable. Having never used gmock myself, I'm leaving this up
to you.

[Deprecated (see juliatuttle), 2012/09/26 20:44:50]

I'm going to leave this, since it works and is not unreasonable.

+}
+
+} // namespace
+
+} // namespace net