styx/dataset/dnstrie/valuetrie_test.go

package dnstrie

import (
	"testing"
)

func TestValueTrie(t *testing.T) {
	t.Run("InsertAndSearchStrings", func(t *testing.T) {
		trie := NewValue[string]()
		domain := "www.example.com"
		value := "192.0.2.1"
		err := trie.Insert(domain, value)
		if err != nil {
			t.Fatalf("Expected no error on insert, got %v", err)
		}

		foundValue, found := trie.Search(domain)
		if !found {
			t.Fatalf("Expected to find domain '%s', but did not", domain)
		}
		if foundValue != value {
			t.Errorf("Expected value '%s', got '%s'", value, foundValue)
		}
	})

	t.Run("SearchNotFound", func(t *testing.T) {
		trie := NewValue[string]()
		err := trie.Insert("example.com", "value")
		if err != nil {
			t.Fatalf("Insert failed: %v", err)
		}

		_, found := trie.Search("nonexistent.com")
		if found {
			t.Error("Expected not to find domain 'nonexistent.com', but did")
		}

		// Search for a path that exists but is not a terminal node
		_, found = trie.Search("com")
		if found {
			t.Error("Expected not to find non-terminal path 'com', but did")
		}
	})

	t.Run("InsertInvalidDomain", func(t *testing.T) {
		trie := NewValue[string]()
		err := trie.Insert("not-a-valid-domain-", "value")
		if err == nil {
			t.Error("Expected an error when inserting an invalid domain, but got nil")
		}
	})

	t.Run("OverwriteValue", func(t *testing.T) {
		trie := NewValue[string]()
		domain := "overwrite.com"
		initialValue := "first"
		overwriteValue := "second"

		err := trie.Insert(domain, initialValue)
		if err != nil {
			t.Fatalf("Initial insert failed: %v", err)
		}
		err = trie.Insert(domain, overwriteValue)
		if err != nil {
			t.Fatalf("Overwrite insert failed: %v", err)
		}

		foundValue, found := trie.Search(domain)
		if !found {
			t.Fatalf("Expected to find domain '%s' after overwrite", domain)
		}
		if foundValue != overwriteValue {
			t.Errorf("Expected overwritten value '%s', got '%s'", overwriteValue, foundValue)
		}
	})

	t.Run("Canonicalization", func(t *testing.T) {
		trie := NewValue[string]()
		value := "canonical"
		// Insert lowercase with trailing dot
		err := trie.Insert("case.example.org.", value)
		if err != nil {
			t.Fatalf("Insert failed: %v", err)
		}

		// Search uppercase without trailing dot
		foundValue, found := trie.Search("CASE.EXAMPLE.ORG")
		if !found {
			t.Fatal("Failed to find domain with different case and no trailing dot")
		}
		if foundValue != value {
			t.Errorf("Expected value '%s' for canonical search, got '%s'", value, foundValue)
		}
	})

	t.Run("InsertAndSearchIntegers", func(t *testing.T) {
		trie := NewValue[int]()
		domain := "int.example.com"
		value := 12345

		err := trie.Insert(domain, value)
		if err != nil {
			t.Fatalf("Expected no error on insert for int trie, got %v", err)
		}

		foundValue, found := trie.Search(domain)
		if !found {
			t.Fatalf("Expected to find domain '%s' in int trie, but did not", domain)
		}
		if foundValue != value {
			t.Errorf("Expected int value %d, got %d", value, foundValue)
		}

		// Search for a non-existent domain in the int trie
		_, found = trie.Search("nonexistent.int.example.com")
		if found {
			t.Error("Found a nonexistent domain in the int trie")
		}
	})

	t.Run("Contains", func(t *testing.T) {
		trie := NewValue[string]()
		domain := "contains.example.com"
		err := trie.Insert(domain, "some-value")
		if err != nil {
			t.Fatalf("Insert failed: %v", err)
		}

		if !trie.Contains(domain) {
			t.Errorf("Expected Contains('%s') to be true, but it was false", domain)
		}

		if trie.Contains("nonexistent." + domain) {
			t.Errorf("Expected Contains for nonexistent domain to be false, but it was true")
		}

		if trie.Contains("example.com") {
			t.Error("Expected Contains for a non-terminal path to be false, but it was true")
		}
	})

	t.Run("MergeTries", func(t *testing.T) {
		trie1 := NewValue[int]()
		trie1.Insert("example.com", 100)
		trie1.Insert("sub.example.com", 200)

		trie2 := NewValue[int]()
		trie2.Insert("google.com", 300)
		trie2.Insert("sub.example.com", 999) // Overlapping domain, new value
		trie2.Insert("another.net", 400)

		trie1.Merge(trie2)

		// Test domains from both tries are present
		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("another.net") {
			t.Error("Merge failed: trie1 should contain 'another.net'")
		}

		// Test that overlapping value was updated from trie2
		val, found := trie1.Search("sub.example.com")
		if !found || val != 999 {
			t.Errorf("Expected value for overlapping domain to be 999, but got %d", val)
		}

		// Test that non-overlapping value from trie1 is intact
		val, found = trie1.Search("example.com")
		if !found || val != 100 {
			t.Errorf("Expected value for 'example.com' to be 100, but got %d", val)
		}

		// Ensure trie2 is not modified
		if _, found := trie2.Search("example.com"); found {
			t.Error("Source trie (trie2) was modified after merge")
		}
	})
}