1 files changed, 48 insertions, 0 deletions
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]);
+}