Land Recent QUIC Changes.

net/tools/quic/quic_dispatcher_test.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/tools/quic/quic_dispatcher.h"
 
 #include <string>
 
 #include "base/strings/string_piece.h"
 #include "net/quic/crypto/crypto_handshake.h"
 #include "net/quic/crypto/quic_crypto_server_config.h"
 #include "net/quic/crypto/quic_random.h"
 #include "net/quic/quic_crypto_stream.h"
 #include "net/quic/test_tools/quic_test_utils.h"
+#include "net/quic/test_tools/quic_test_writer.h"
 #include "net/tools/epoll_server/epoll_server.h"
 #include "net/tools/quic/quic_time_wait_list_manager.h"
 #include "net/tools/quic/test_tools/quic_dispatcher_peer.h"
 #include "net/tools/quic/test_tools/quic_test_utils.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 using base::StringPiece;
 using net::EpollServer;
 using net::test::MockSession;
+using net::test::QuicTestWriter;
 using net::tools::test::MockConnection;
 using std::make_pair;
 using testing::_;
 using testing::DoAll;
 using testing::Invoke;
 using testing::InSequence;
 using testing::Return;
 using testing::WithoutArgs;
 
 namespace net {
@@ skipping 228 matching lines @@
   IPEndPoint addr(net::test::Loopback4(), 1);
   QuicGuid guid = 1;
   // Dispatcher forwards all packets for this guid to the time wait list
   // manager.
   EXPECT_CALL(dispatcher_, CreateQuicSession(_, _, _)).Times(0);
   EXPECT_CALL(*time_wait_list_manager, ProcessPacket(_, _, guid, _)).Times(1);
   string data = "foo";
   ProcessPacket(addr, guid, false, "foo");
 }
 
+class BlockingWriter : public QuicTestWriter {
+ public:
+  BlockingWriter() : write_blocked_(false) {}
+
+  virtual bool IsWriteBlocked() const OVERRIDE { return write_blocked_; }
+  virtual void SetWritable() OVERRIDE { write_blocked_ = false; }
+  virtual bool IsWriteBlockedDataBuffered() const OVERRIDE { return false; }
+
+  virtual WriteResult WritePacket(
+      const char* buffer, size_t buf_len, const IPAddressNumber& self_address,
+      const IPEndPoint& peer_address,
+      QuicBlockedWriterInterface* blocked_writer) OVERRIDE {
+    if (write_blocked_) {
+      return WriteResult(WRITE_STATUS_BLOCKED, EAGAIN);
+    } else {
+      return writer()->WritePacket(buffer, buf_len, self_address, peer_address,
+                                   blocked_writer);
+    }
+  }
+
+  bool write_blocked_;
+};
+
 class QuicWriteBlockedListTest : public QuicDispatcherTest {
  public:
   virtual void SetUp() {
+    writer_ = new BlockingWriter;
+    QuicDispatcherPeer::UseWriter(&dispatcher_, writer_);
+
     IPEndPoint addr(net::test::Loopback4(), 1);
 
     EXPECT_CALL(dispatcher_, CreateQuicSession(_, _, addr))
         .WillOnce(testing::Return(CreateSession(
                       &dispatcher_, 1, addr, &session1_)));
     ProcessPacket(addr, 1, true, "foo");
 
     EXPECT_CALL(dispatcher_, CreateQuicSession(_, _, addr))
         .WillOnce(testing::Return(CreateSession(
                       &dispatcher_, 2, addr, &session2_)));
     ProcessPacket(addr, 2, true, "bar");
 
     blocked_list_ = dispatcher_.write_blocked_list();
   }
 
   virtual void TearDown() {
     EXPECT_CALL(*connection1(), SendConnectionClose(QUIC_PEER_GOING_AWAY));
     EXPECT_CALL(*connection2(), SendConnectionClose(QUIC_PEER_GOING_AWAY));
     dispatcher_.Shutdown();
   }
 
   bool SetBlocked() {
-    QuicDispatcherPeer::SetWriteBlocked(&dispatcher_);
+    writer_->write_blocked_ = true;;
     return true;
   }
 
  protected:
+  BlockingWriter* writer_;
   QuicDispatcher::WriteBlockedList* blocked_list_;
 };
 
 TEST_F(QuicWriteBlockedListTest, BasicOnCanWrite) {
   // No OnCanWrite calls because no connections are blocked.
   dispatcher_.OnCanWrite();
 
   // Register connection 1 for events, and make sure it's nofitied.
   blocked_list_->insert(make_pair(connection1(), true));
   EXPECT_CALL(*connection1(), OnCanWrite());
@@ skipping 100 matching lines @@
 
   // And we'll resume where we left off when we get another call.
   EXPECT_CALL(*connection2(), OnCanWrite());
   dispatcher_.OnCanWrite();
 }
 
 }  // namespace
 }  // namespace test
 }  // namespace tools
 }  // namespace net