Add unittest for DiscretePDF

3 files changed, 63 insertions, 0 deletions

diff --git a/unittests/CMakeLists.txt b/unittests/CMakeLists.txt
index 63934336..00d4cbfa 100644
--- a/unittests/CMakeLists.txt
+++ b/unittests/CMakeLists.txt
@@ -85,6 +85,7 @@ add_subdirectory(Expr)
 add_subdirectory(Ref)
 add_subdirectory(Solver)
 add_subdirectory(TreeStream)
+add_subdirectory(DiscretePDF)
 
 # Set up lit configuration
 set (UNIT_TEST_EXE_SUFFIX "Test")
diff --git a/unittests/DiscretePDF/CMakeLists.txt b/unittests/DiscretePDF/CMakeLists.txt
new file mode 100644
index 00000000..07ec1f89
--- /dev/null
+++ b/unittests/DiscretePDF/CMakeLists.txt
@@ -0,0 +1,4 @@
+add_klee_unit_test(DiscretePDFTest
+  DiscretePDFTest.cpp)
+# FIXME add the following line to link against libgtest.a
+target_link_libraries(DiscretePDFTest PRIVATE kleaverSolver)
\ No newline at end of file
diff --git a/unittests/DiscretePDF/DiscretePDFTest.cpp b/unittests/DiscretePDF/DiscretePDFTest.cpp
new file mode 100644
index 00000000..2180360c
--- /dev/null
+++ b/unittests/DiscretePDF/DiscretePDFTest.cpp
@@ -0,0 +1,58 @@
+#include "klee/Internal/ADT/DiscretePDF.h"
+#include "gtest/gtest.h"
+#include <iostream>
+#include <vector>
+
+int finished = 0;
+
+using namespace klee;
+
+TEST(DiscretePDFTest, Rotatation) {
+  DiscretePDF<int> testTree;
+
+  ASSERT_TRUE(testTree.empty());
+
+  for (auto i = 0; i < 9; ++i)
+    testTree.insert(i + 10, 0.1 * i);
+
+  for (auto i = 9; i > 0; --i)
+    testTree.insert(i + 100, 0.01 * i);
+
+  ASSERT_FALSE(testTree.empty());
+
+  ASSERT_TRUE(testTree.inTree(101));
+  testTree.remove(101);
+  ASSERT_FALSE(testTree.inTree(101));
+  testTree.insert(101, 0.01);
+  ASSERT_EQ(0.01, testTree.getWeight(101));
+
+  ASSERT_EQ(11, testTree.choose(0.01));
+  testTree.choose(0);
+  testTree.choose(0.9999999);
+
+  testTree.update(101, 0.9);
+
+  for (auto i = 0; i < 9; ++i)
+    testTree.remove(i + 10);
+  for (auto i = 0; i < 50; ++i) {
+    testTree.insert(50 + i, 2 * i);
+  }
+
+  for (auto i = 9; i > 0; --i)
+    testTree.remove(i + 100);
+
+  testTree.insert(1, 0);
+#ifndef NDEBUG
+  ASSERT_DEATH({ testTree.insert(1, 0); }, "already in tree");
+#endif
+
+  while (!testTree.empty()) {
+    testTree.remove(testTree.choose(0));
+  }
+
+  testTree.insert(1, 1);
+  testTree.insert(2, 2);
+
+  ASSERT_EQ(1, testTree.getWeight(1));
+  ASSERT_EQ(2, testTree.getWeight(2));
+}