Better trie implementations

dataset/dnstrie/trie_test.go

@@ -0,0 +1,117 @@
+package dnstrie
+
+import "testing"
+
+func TestTrie(t *testing.T) {
+	t.Run("InsertAndContains", func(t *testing.T) {
+		trie := New()
+		domain := "www.example.com"
+		err := trie.Insert(domain)
+		if err != nil {
+			t.Fatalf("Expected no error on insert, got %v", err)
+		}
+
+		if !trie.Contains(domain) {
+			t.Errorf("Expected Contains('%s') to be true, but it was false", domain)
+		}
+	})
+
+	t.Run("ContainsNotFound", func(t *testing.T) {
+		trie := New()
+		err := trie.Insert("example.com")
+		if err != nil {
+			t.Fatalf("Insert failed: %v", err)
+		}
+
+		if trie.Contains("nonexistent.com") {
+			t.Error("Expected not to find domain 'nonexistent.com', but did")
+		}
+
+		// Check for a path that exists but is not a terminal node
+		if trie.Contains("com") {
+			t.Error("Expected not to find non-terminal path 'com', but did")
+		}
+	})
+
+	t.Run("InsertInvalidDomain", func(t *testing.T) {
+		trie := New()
+		err := trie.Insert("not-a-valid-domain-")
+		if err == nil {
+			t.Error("Expected an error when inserting an invalid domain, but got nil")
+		}
+	})
+
+	t.Run("Canonicalization", func(t *testing.T) {
+		trie := New()
+		// Insert lowercase with trailing dot
+		err := trie.Insert("case.example.org.")
+		if err != nil {
+			t.Fatalf("Insert failed: %v", err)
+		}
+
+		// Check contains with uppercase without trailing dot
+		if !trie.Contains("CASE.EXAMPLE.ORG") {
+			t.Fatal("Failed to find domain with different case and no trailing dot")
+		}
+	})
+
+	t.Run("MultipleInsertions", func(t *testing.T) {
+		trie := New()
+		domains := []string{
+			"example.com",
+			"www.example.com",
+			"api.example.com",
+			"google.com",
+		}
+		for _, domain := range domains {
+			if err := trie.Insert(domain); err != nil {
+				t.Fatalf("Insert failed for %s: %v", domain, err)
+			}
+		}
+
+		for _, domain := range domains {
+			if !trie.Contains(domain) {
+				t.Errorf("Expected to find %s, but did not", domain)
+			}
+		}
+
+		if trie.Contains("ftp.example.com") {
+			t.Error("Found domain 'ftp.example.com' which was not inserted")
+		}
+	})
+
+	t.Run("MergeTries", func(t *testing.T) {
+		trie1 := New()
+		trie1.Insert("example.com")
+		trie1.Insert("sub.example.com")
+
+		trie2 := New()
+		trie2.Insert("google.com")
+		trie2.Insert("sub.example.com") // Overlapping domain
+		trie2.Insert("another.net")
+
+		trie1.Merge(trie2)
+
+		// Test domains from both tries
+		if !trie1.Contains("example.com") {
+			t.Error("Merge failed: trie1 should contain 'example.com'")
+		}
+		if !trie1.Contains("google.com") {
+			t.Error("Merge failed: trie1 should contain 'google.com'")
+		}
+		if !trie1.Contains("sub.example.com") {
+			t.Error("Merge failed: trie1 should contain overlapping 'sub.example.com'")
+		}
+		if !trie1.Contains("another.net") {
+			t.Error("Merge failed: trie1 should contain 'another.net'")
+		}
+
+		// Ensure trie2 is not modified
+		if !trie2.Contains("google.com") {
+			t.Error("Source trie (trie2) should not be modified after merge")
+		}
+		if trie2.Contains("example.com") {
+			t.Error("Source trie (trie2) was modified after merge")
+		}
+	})
+}