3 files changed, 62 insertions, 0 deletions
diff --git a/unittests/TreeStream/CMakeLists.txt b/unittests/TreeStream/CMakeLists.txt
new file mode 100644
index 00000000..4b298724
--- /dev/null
+++ b/unittests/TreeStream/CMakeLists.txt
@@ -0,0 +1,3 @@
+add_klee_unit_test(TreeStreamTest
+  TreeStreamTest.cpp)
+target_link_libraries(TreeStreamTest PRIVATE kleeBasic)
diff --git a/unittests/TreeStream/Makefile b/unittests/TreeStream/Makefile
new file mode 100644
index 00000000..a56659b9
--- /dev/null
+++ b/unittests/TreeStream/Makefile
@@ -0,0 +1,11 @@
+##===- unittests/Expr/Makefile -----------------------------*- Makefile -*-===##
+
+LEVEL := ../..
+include $(LEVEL)/Makefile.config
+
+TESTNAME := TreeStreamTest
+USEDLIBS := kleeBasic.a kleeSupport.a
+LINK_COMPONENTS := support
+
+include $(LLVM_SRC_ROOT)/unittests/Makefile.unittest
+
diff --git a/unittests/TreeStream/TreeStreamTest.cpp b/unittests/TreeStream/TreeStreamTest.cpp
new file mode 100644
index 00000000..09d62c6e
--- /dev/null
+++ b/unittests/TreeStream/TreeStreamTest.cpp
@@ -0,0 +1,48 @@
+#include "klee/Internal/ADT/TreeStream.h"
+#include <vector>
+#include <cstring>
+
+#include "gtest/gtest.h"
+
+using namespace klee;
+
+/* Basic test, checking that after writing "abc" and then "defg", we
+   get a {'a', 'b', 'c', 'c', 'd', 'e', 'f', 'g' } back.  */
+TEST(TreeStreamTest, Basic) {
+  TreeStreamWriter tsw("tsw1.out");
+  ASSERT_TRUE(tsw.good());
+  
+  TreeOStream tos = tsw.open();
+  tos.write("abc", 3);
+  tos.write("defg", 4);
+  tos.flush();
+
+  std::vector<unsigned char> out;
+  tsw.readStream(tos.getID(), out);
+  ASSERT_EQ(out.size(), 7);
+  
+  for (unsigned char c = 'a'; c <= 'g'; c++)
+    ASSERT_EQ(out[c - 'a'], c);
+}
+
+
+/* This tests the case when we perform a write with a size larger than
+   the buffer size, which is a constant set to 4*4096.  This test fails 
+   without #704 */
+TEST(TreeStreamTest, WriteLargerThanBufferSize) {
+  TreeStreamWriter tsw("tsw2.out");
+  ASSERT_TRUE(tsw.good());
+  
+  TreeOStream tos = tsw.open();
+#define NBYTES 5*4096
+  char buf[NBYTES];
+  memset(buf, 'A', sizeof(buf));
+  tos.write(buf, NBYTES);
+  tos.flush();
+
+  std::vector<unsigned char> out;
+  tsw.readStream(tos.getID(), out);
+  ASSERT_EQ(out.size(), NBYTES);
+  for (unsigned i=0; i<out.size(); i++)
+    ASSERT_EQ('A', out[i]);
+}