Patch set 2 of Daniel Ziegler's CL

net/quic/test_tools/quic_test_utils.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/quic/test_tools/quic_test_utils.h"
 
+#include "base/memory/singleton.h"
 #include "base/sha1.h"
 #include "base/stl_util.h"
 #include "base/strings/string_number_conversions.h"
 #include "net/quic/crypto/crypto_framer.h"
 #include "net/quic/crypto/crypto_handshake.h"
 #include "net/quic/crypto/crypto_utils.h"
 #include "net/quic/crypto/null_encrypter.h"
 #include "net/quic/crypto/quic_decrypter.h"
 #include "net/quic/crypto/quic_encrypter.h"
 #include "net/quic/quic_framer.h"
@@ skipping 197 matching lines @@
 }
 
 QuicAlarm* MockHelper::CreateAlarm(QuicAlarm::Delegate* delegate) {
   return new TestAlarm(delegate);
 }
 
 void MockHelper::AdvanceTime(QuicTime::Delta delta) {
   clock_.AdvanceTime(delta);
 }
 
+namespace {
+class NiceMockPacketWriterFactory
+    : public QuicConnection::PacketWriterFactory {
+ public:
+  virtual ~NiceMockPacketWriterFactory() {}
+
+  static NiceMockPacketWriterFactory* GetInstance() {
+    return Singleton<NiceMockPacketWriterFactory>::get();
+  }
+
+  virtual QuicPacketWriter* Create(
+      QuicConnection* /*connection*/) const override {
+    return new testing::NiceMock<MockPacketWriter>();
+  }
+
+ private:
+  friend struct DefaultSingletonTraits<NiceMockPacketWriterFactory>;
+  NiceMockPacketWriterFactory() {}
+  DISALLOW_COPY_AND_ASSIGN(NiceMockPacketWriterFactory);
+};
+}  // namespace
+
 MockConnection::MockConnection(bool is_server)
     : QuicConnection(kTestConnectionId,
                      IPEndPoint(TestPeerIPAddress(), kTestPort),
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     *NiceMockPacketWriterFactory::GetInstance(),
+                     /* owns_writer= */ true,
                      is_server, QuicSupportedVersions()),
       helper_(helper()) {
 }
 
 MockConnection::MockConnection(IPEndPoint address,
                                bool is_server)
     : QuicConnection(kTestConnectionId, address,
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     *NiceMockPacketWriterFactory::GetInstance(),
+                     /* owns_writer= */ true,
                      is_server, QuicSupportedVersions()),
       helper_(helper()) {
 }
 
 MockConnection::MockConnection(QuicConnectionId connection_id,
                                bool is_server)
     : QuicConnection(connection_id,
                      IPEndPoint(TestPeerIPAddress(), kTestPort),
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     *NiceMockPacketWriterFactory::GetInstance(),
+                     /* owns_writer= */ true,
                      is_server, QuicSupportedVersions()),
       helper_(helper()) {
 }
 
 MockConnection::MockConnection(bool is_server,
                                const QuicVersionVector& supported_versions)
     : QuicConnection(kTestConnectionId,
                      IPEndPoint(TestPeerIPAddress(), kTestPort),
                      new testing::NiceMock<MockHelper>(),
-                     new testing::NiceMock<MockPacketWriter>(),
-                     true  /* owns_writer */,
+                     *NiceMockPacketWriterFactory::GetInstance(),
+                     /* owns_writer= */ true,
                      is_server, supported_versions),
       helper_(helper()) {
 }
 
 MockConnection::~MockConnection() {
 }
 
 void MockConnection::AdvanceTime(QuicTime::Delta delta) {
   static_cast<MockHelper*>(helper())->AdvanceTime(delta);
 }
@@ skipping 327 matching lines @@
       kInitialSessionFlowControlWindowForTest);
   return config;
 }
 
 QuicVersionVector SupportedVersions(QuicVersion version) {
   QuicVersionVector versions;
   versions.push_back(version);
   return versions;
 }
 
+TestWriterFactory::TestWriterFactory() : current_writer_(NULL) {}
+TestWriterFactory::~TestWriterFactory() {}
+
+QuicPacketWriter* TestWriterFactory::Create(QuicServerPacketWriter* writer,
+                                            QuicConnection* connection) {
+  return new PerConnectionPacketWriter(this, writer, connection);
+}
+
+void TestWriterFactory::OnPacketSent(WriteResult result) {
+  if (current_writer_ != NULL) {
+    current_writer_->connection()->OnPacketSent(result);
+    current_writer_ = NULL;
+  }
+}
+
+void TestWriterFactory::Unregister(PerConnectionPacketWriter* writer) {
+  if (current_writer_ == writer) {
+    current_writer_ = NULL;
+  }
+}
+
+TestWriterFactory::PerConnectionPacketWriter::PerConnectionPacketWriter(
+    TestWriterFactory* factory,
+    QuicServerPacketWriter* writer,
+    QuicConnection* connection)
+    : QuicPerConnectionPacketWriter(writer, connection),
+      factory_(factory) {
+}
+
+TestWriterFactory::PerConnectionPacketWriter::~PerConnectionPacketWriter() {
+  factory_->Unregister(this);
+}
+
+WriteResult TestWriterFactory::PerConnectionPacketWriter::WritePacket(
+    const char* buffer,
+    size_t buf_len,
+    const IPAddressNumber& self_address,
+    const IPEndPoint& peer_address) {
+  // A DCHECK(factory_current_writer_ == NULL) would be wrong here -- this class
+  // may be used in a setting where connection()->OnPacketSent() is called in a
+  // different way, so TestWriterFactory::OnPacketSent might never be called.
+  factory_->current_writer_ = this;
+  return QuicPerConnectionPacketWriter::WritePacket(buffer,
+                                                    buf_len,
+                                                    self_address,
+                                                    peer_address);
+}
+
 }  // namespace test
 }  // namespace net