Keep pool of pre-connected DNS sockets

net/dns/dns_session_unittest.cc

+// 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 "base/bind.h"
+#include "base/memory/scoped_ptr.h"
+#include "base/rand_util.h"
+#include "net/dns/dns_protocol.h"
+#include "net/socket/socket_test_util.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace net {
+
+namespace {
+
+class TestSocketFactory;
+
+class TestUDPClientSocket : public MockUDPClientSocket {
+ public:
+  TestUDPClientSocket(TestSocketFactory* factory,
+                      SocketDataProvider* data,
+                      net::NetLog* net_log)
+      : MockUDPClientSocket(data, net_log), factory_(factory) {}
+  virtual ~TestUDPClientSocket();
+ private:
+  TestSocketFactory* factory_;
+};
+
+// TODO(ttuttle): Does it make sense to subclass MockClientSocketFactory?
+//                We're basically entirely overriding the functionality it
+//                provides.
+class TestSocketFactory : public MockClientSocketFactory {
+ public:
+  TestSocketFactory() : pending_created_(0), pending_destroyed_(0) {}
+  virtual ~TestSocketFactory() {}
+
+  virtual DatagramClientSocket* CreateDatagramClientSocket(
+      DatagramSocket::BindType bind_type,
+      const RandIntCallback& rand_int_cb,
+      net::NetLog* net_log,
+      const net::NetLog::Source& source) OVERRIDE {
+    // 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();
+    TestUDPClientSocket* socket = new TestUDPClientSocket(this,
+                                                          data_provider,
+                                                          net_log);
+    data_provider->set_socket(socket);
+    pending_created_++;
+    return socket;
+  }
+
+  void OnSocketDestroyed() {
+    ++pending_destroyed_;
+  }
+
+  void ExpectSocketsCreated(int created) {
+    EXPECT_EQ(created, pending_created_);
+    pending_created_ = 0;
+  }
+
+  void ExpectSocketsDestroyed(int destroyed) {
+    EXPECT_EQ(destroyed, pending_destroyed_);
+    pending_destroyed_ = 0;
+  }
+
+ protected:
+  int pending_created_;
+  int pending_destroyed_;
+};
+
+TestUDPClientSocket::~TestUDPClientSocket() {
+  factory_->OnSocketDestroyed();
+}
+
+typedef DnsSession::SocketLease SocketLease;
+static const size_t kBufferSize = 4;
+
+class DnsSessionTest : public testing::Test {
+ public:
+  void Initialize() {
+    InitializeFactory();
+    InitializeConfig(1);
+    InitializeSession();
+  }
+
+ protected:
+  DnsConfig config_;
+  scoped_ptr<TestSocketFactory> socket_factory_;
+  scoped_refptr<DnsSession> session_;
+
+  void InitializeFactory() {
+    socket_factory_.reset(new TestSocketFactory());
+  }
+
+  // TODO(ttuttle): Refactor this and the one in DnsTransactionTest somewhere.
+  void InitializeConfig(unsigned num_servers) {
+    CHECK_LE(num_servers, 255u);
+    config_.nameservers.clear();
+    IPAddressNumber dns_ip;
+    bool rv = ParseIPLiteralToNumber("192.168.1.0", &dns_ip);
+    EXPECT_TRUE(rv);
+    for (unsigned i = 0; i < num_servers; ++i) {
+      dns_ip[3] = i;
+      IPEndPoint dns_endpoint(dns_ip, dns_protocol::kDefaultPort);
+      config_.nameservers.push_back(dns_endpoint);
+    }
+  }
+
+  void InitializeSession() {
+    session_ = new DnsSession(config_,
+                              socket_factory_.get(),
+                              base::Bind(&base::RandInt),
+                              NULL /* NetLog */);
+  }
+
+  void ExpectCreated(int num) {
+    socket_factory_->ExpectSocketsCreated(num);

[szym, 2012/09/10 19:31:00]

Wrapping EXPECT_* in methods makes it more difficult to track down where the
test fails, because the line printed will be from here, not from where the
Expect* wrapper is called. If possible, have it return a boolean and EXPECT_* it
in the test body.

[Deprecated (see juliatuttle), 2012/09/14 15:28:18]

Done.

+  }
+
+  void ExpectDestroyed(int num) {
+    socket_factory_->ExpectSocketsDestroyed(num);
+  }
+};
+
+/*
+ * TODO(ttuttle):
+ * Mock out SocketPools and make sure the pools get the right calls.
+ * Test the default SocketPool implementation.
+ *   ...while in DnsSession or separately?
+ */
+
+}
+
+}