1 files changed, 56 insertions, 15 deletions

diff --git a/test.cc b/test.cc
index 5105f09..2c38a3f 100644
--- a/test.cc
+++ b/test.cc
@@ -1,24 +1,65 @@
 #include "trie.h"
 
-#include <cassert>
+#include "gtest/gtest.h"
+
 #include <cstdint>
 
-int main() {
-	Trie<uint8_t> test;
+class TrieTest : public ::testing::Test { };
+
+TEST_F(TrieTest, Basic) {
+	Trie<uint8_t> trie;
+
+	trie.add({1, 1, 1, 1});
+	trie.add({1, 2, 1, 2});
+	trie.add({2, 3, 4, 5});
+	trie.add({2, 3, 1, 2});
+
+	EXPECT_EQ(trie.resolve({1, 1, 1, 1}).first,  true);
+	EXPECT_NE(trie.resolve({1, 1, 1, 1}).second, nullptr);
+
+	EXPECT_EQ(trie.resolve({1, 2, 1, 2}).first,  true);
+	EXPECT_NE(trie.resolve({1, 2, 1, 2}).second, nullptr);
+
+	EXPECT_EQ(trie.resolve({1, 2}).first,        true);
+	EXPECT_NE(trie.resolve({1, 2}).second,       nullptr);
+
+	EXPECT_EQ(trie.resolve({1, 1, 2, 3}).first,  false);
+	EXPECT_EQ(trie.resolve({1, 1, 2, 3}).second, nullptr);
+
+	EXPECT_EQ(trie.resolve({2, 1, 4}).first,     false);
+	EXPECT_EQ(trie.resolve({2, 1, 4}).second,    nullptr);
+}
 
-	test.add({1, 2, 3});
-	test.add({1, 2, 4});
-	test.add({2, 1});
-	test.add({2, 1, 1});
+TEST_F(TrieTest, Remove) {
+	Trie<uint8_t> trie;
 
-	assert(test.resolve({1, 2}).second    != nullptr);
-	assert(test.resolve({1, 2, 3}).second != nullptr);
-	assert(test.resolve({1, 2, 4}).second != nullptr);
-	assert(test.resolve({3}).second       == nullptr);
+	trie.add({1, 1, 1, 1});
+	trie.add({1, 2, 1, 2});
+	trie.add({2, 3, 4, 5});
+	trie.add({2, 3, 1, 2});
+
+	trie.remove({1, 1, 1, 1});
+
+	EXPECT_EQ(trie.resolve({1, 1, 1, 1}).first,  false);
+	EXPECT_EQ(trie.resolve({1, 1, 1, 1}).second, nullptr);
+
+	EXPECT_EQ(trie.resolve({1, 2, 1, 2}).first,  true);
+	EXPECT_NE(trie.resolve({1, 2, 1, 2}).second, nullptr);
+
+	trie.remove({2, 3});
+
+	EXPECT_EQ(trie.resolve({2, 3, 4, 5}).first,  false);
+	EXPECT_EQ(trie.resolve({2, 3, 4, 5}).second, nullptr);
+
+	EXPECT_EQ(trie.resolve({2, 3}).first,        false);
+	EXPECT_EQ(trie.resolve({2, 3}).second,       nullptr);
+
+	EXPECT_EQ(trie.resolve({2}).first,           true);
+	EXPECT_NE(trie.resolve({2}).second,          nullptr);
+}
 
-	test.remove({1, 2});
+int main(int argc, char **argv) {
+	testing::InitGoogleTest(&argc, argv);
 
-	assert(test.resolve({1, 2, 4}).second == nullptr);
-	assert(test.resolve({1, 2, 3}).second == nullptr);
-	assert(test.resolve({1, 2}).second    == nullptr);
+	return RUN_ALL_TESTS();
 }