Add trimAtMost functions

[folly.git] / folly / io / test / IOBufQueueTest.cpp

diff --git a/folly/io/test/IOBufQueueTest.cpp b/folly/io/test/IOBufQueueTest.cpp
index 8ba3d7b47c85aca0c8b4aed192ad7d6fd1333478..8050812e229fba23148b95a7b92e955c22a4d156 100644 (file)
--- a/folly/io/test/IOBufQueueTest.cpp
+++ b/folly/io/test/IOBufQueueTest.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,9 +20,8 @@
 #include <stdexcept>
 #include <string.h>
 
-#include <gtest/gtest.h>
-
 #include <folly/Range.h>
+#include <folly/portability/GTest.h>
 
 using folly::IOBuf;
 using folly::IOBufQueue;
@@ -44,8 +43,7 @@ struct Initializer {
 };
 Initializer initializer;
 
-unique_ptr<IOBuf>
-stringToIOBuf(const char* s, uint32_t len) {
+unique_ptr<IOBuf> stringToIOBuf(const char* s, size_t len) {
   unique_ptr<IOBuf> buf = IOBuf::create(len);
   memcpy(buf->writableTail(), s, len);
   buf->append(len);
@@ -160,6 +158,15 @@ TEST(IOBufQueue, Split) {
   checkConsistency(queue);
 }
 
+TEST(IOBufQueue, SplitAtMost) {
+  IOBufQueue queue(clOptions);
+  queue.append(stringToIOBuf(SCL("Hello,")));
+  queue.append(stringToIOBuf(SCL(" World")));
+  auto buf = queue.splitAtMost(9999);
+  EXPECT_EQ(buf->computeChainDataLength(), 12);
+  EXPECT_TRUE(queue.empty());
+}
+
 TEST(IOBufQueue, Preallocate) {
   IOBufQueue queue(clOptions);
   queue.append(string("Hello"));
@@ -261,6 +268,76 @@ TEST(IOBufQueue, Trim) {
   checkConsistency(queue);
 }
 
+TEST(IOBufQueue, TrimStartAtMost) {
+  IOBufQueue queue(clOptions);
+  unique_ptr<IOBuf> a = IOBuf::create(4);
+  a->append(4);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(6);
+  a->append(6);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(8);
+  a->append(8);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(10);
+  a->append(10);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(28, queue.front()->computeChainDataLength());
+  EXPECT_EQ(4, queue.front()->length());
+
+  queue.trimStartAtMost(1);
+  checkConsistency(queue);
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(27, queue.front()->computeChainDataLength());
+  EXPECT_EQ(3, queue.front()->length());
+
+  queue.trimStartAtMost(50);
+  checkConsistency(queue);
+  EXPECT_EQ(nullptr, queue.front());
+  EXPECT_EQ(0, queue.chainLength());
+}
+
+TEST(IOBufQueue, TrimEndAtMost) {
+  IOBufQueue queue(clOptions);
+  unique_ptr<IOBuf> a = IOBuf::create(4);
+  a->append(4);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(6);
+  a->append(6);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(8);
+  a->append(8);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+  a = IOBuf::create(10);
+  a->append(10);
+  queue.append(std::move(a));
+  checkConsistency(queue);
+
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(28, queue.front()->computeChainDataLength());
+  EXPECT_EQ(4, queue.front()->length());
+
+  queue.trimEndAtMost(1);
+  checkConsistency(queue);
+  EXPECT_EQ(4, queue.front()->countChainElements());
+  EXPECT_EQ(27, queue.front()->computeChainDataLength());
+  EXPECT_EQ(4, queue.front()->length());
+
+  queue.trimEndAtMost(50);
+  checkConsistency(queue);
+  EXPECT_EQ(nullptr, queue.front());
+  EXPECT_EQ(0, queue.chainLength());
+}
+
 TEST(IOBufQueue, TrimPack) {
   IOBufQueue queue(clOptions);
   unique_ptr<IOBuf> a = IOBuf::create(64);
@@ -385,3 +462,19 @@ TEST(IOBufQueue, AppendToString) {
   queue.appendToString(s);
   EXPECT_EQ("hello world", s);
 }
+
+TEST(IOBufQueue, Gather) {
+  IOBufQueue queue;
+
+  queue.append(stringToIOBuf(SCL("hello ")));
+  queue.append(stringToIOBuf(SCL("world")));
+
+  EXPECT_EQ(queue.front()->length(), 6);
+  queue.gather(11);
+  EXPECT_EQ(queue.front()->length(), 11);
+
+  StringPiece s(
+      reinterpret_cast<const char*>(queue.front()->data()),
+      queue.front()->length());
+  EXPECT_EQ("hello world", s);
+}