diff --git a/quic/api/test/IoBufQuicBatchTest.cpp b/quic/api/test/IoBufQuicBatchTest.cpp index f1c81021b..18f18339c 100644 --- a/quic/api/test/IoBufQuicBatchTest.cpp +++ b/quic/api/test/IoBufQuicBatchTest.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -18,45 +19,11 @@ constexpr const auto kMaxBufs = 10; namespace quic { namespace testing { -class TestPacketBatchWriter : public IOBufBatchWriter { - public: - explicit TestPacketBatchWriter(int maxBufs) : maxBufs_(maxBufs) {} - ~TestPacketBatchWriter() override { - CHECK_EQ(bufNum_, 0); - CHECK_EQ(bufSize_, 0); - } - - void reset() override { - bufNum_ = 0; - bufSize_ = 0; - } - - bool append( - std::unique_ptr&& /*unused*/, - size_t size, - const folly::SocketAddress& /*unused*/, - folly::AsyncUDPSocket* /*unused*/) override { - bufNum_++; - bufSize_ += size; - return ((maxBufs_ < 0) || (bufNum_ >= maxBufs_)); - } - ssize_t write( - folly::AsyncUDPSocket& /*unused*/, - const folly::SocketAddress& /*unused*/) override { - return bufSize_; - } - - private: - int maxBufs_{0}; - int bufNum_{0}; - size_t bufSize_{0}; -}; - void RunTest(int numBatch) { folly::EventBase evb; folly::AsyncUDPSocket sock(&evb); - auto batchWriter = BatchWriterPtr(new TestPacketBatchWriter(numBatch)); + auto batchWriter = BatchWriterPtr(new test::TestPacketBatchWriter(numBatch)); folly::SocketAddress peerAddress{"127.0.0.1", 1234}; QuicClientConnectionState conn( FizzClientQuicHandshakeContext::Builder().build()); diff --git a/quic/common/test/TestUtils.cpp b/quic/common/test/TestUtils.cpp index 894f01c38..225b8f423 100644 --- a/quic/common/test/TestUtils.cpp +++ b/quic/common/test/TestUtils.cpp @@ -791,5 +791,26 @@ void FizzCryptoTestFactory::setDefault() { return FizzCryptoFactory::makePacketNumberCipher(secret); })); } + +void TestPacketBatchWriter::reset() { + bufNum_ = 0; + bufSize_ = 0; +} + +bool TestPacketBatchWriter::append( + std::unique_ptr&& /*unused*/, + size_t size, + const folly::SocketAddress& /*unused*/, + folly::AsyncUDPSocket* /*unused*/) { + bufNum_++; + bufSize_ += size; + return ((maxBufs_ < 0) || (bufNum_ >= maxBufs_)); +} + +ssize_t TestPacketBatchWriter::write( + folly::AsyncUDPSocket& /*unused*/, + const folly::SocketAddress& /*unused*/) { + return bufSize_; +} } // namespace test } // namespace quic diff --git a/quic/common/test/TestUtils.h b/quic/common/test/TestUtils.h index 7fed2f35f..a78404a64 100644 --- a/quic/common/test/TestUtils.h +++ b/quic/common/test/TestUtils.h @@ -8,6 +8,7 @@ #pragma once +#include #include #include #include @@ -350,5 +351,31 @@ class FizzCryptoTestFactory : public FizzCryptoFactory { mutable std::unique_ptr packetNumberCipher_; }; +class TestPacketBatchWriter : public IOBufBatchWriter { + public: + explicit TestPacketBatchWriter(int maxBufs) : maxBufs_(maxBufs) {} + ~TestPacketBatchWriter() override { + CHECK_EQ(bufNum_, 0); + CHECK_EQ(bufSize_, 0); + } + + void reset() override; + + bool append( + std::unique_ptr&& /*unused*/, + size_t size, + const folly::SocketAddress& /*unused*/, + folly::AsyncUDPSocket* /*unused*/); + + ssize_t write( + folly::AsyncUDPSocket& /*unused*/, + const folly::SocketAddress& /*unused*/); + + private: + int maxBufs_{0}; + int bufNum_{0}; + size_t bufSize_{0}; +}; + } // namespace test } // namespace quic