118 lines
2.8 KiB
Go
118 lines
2.8 KiB
Go
package embedding
|
|
|
|
import (
|
|
"reflect"
|
|
"testing"
|
|
)
|
|
|
|
func TestIndex_AddAndSearch_SingleTerm(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ix := NewIndex[int]()
|
|
ix.Add("go is a programming language", 1)
|
|
ix.Add("rust is a systems language", 2)
|
|
ix.Add("python is great", 3)
|
|
|
|
hits := ix.Search("language")
|
|
if len(hits) != 2 {
|
|
t.Fatalf("Search(\"language\") returned %d hits, want 2", len(hits))
|
|
}
|
|
if got := []int{hits[0].ID, hits[1].ID}; !reflect.DeepEqual(got, []int{1, 2}) {
|
|
t.Errorf("top hits = %v, want [1 2]", got)
|
|
}
|
|
}
|
|
|
|
func TestIndex_SearchRelevanceOrdersByFrequency(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ix := NewIndex[int]()
|
|
ix.Add("cat cat cat dog", 1)
|
|
ix.Add("dog dog cat", 2)
|
|
|
|
hits := ix.Search("cat")
|
|
if len(hits) != 2 {
|
|
t.Fatalf("Search(\"cat\") returned %d hits, want 2", len(hits))
|
|
}
|
|
// Doc 1 has cat 3 times, doc 2 once, so doc 1 ranks higher.
|
|
if hits[0].ID != 1 {
|
|
t.Errorf("top hit = %d, want 1 (higher tf)", hits[0].ID)
|
|
}
|
|
if hits[1].ID != 2 {
|
|
t.Errorf("second hit = %d, want 2", hits[1].ID)
|
|
}
|
|
}
|
|
|
|
func TestIndex_SearchOverlapsMultipleTerms(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ix := NewIndex[int]()
|
|
ix.Add("the red fox", 1)
|
|
ix.Add("the blue whale", 2)
|
|
ix.Add("a red car", 3)
|
|
|
|
hits := ix.Search("red fox")
|
|
if len(hits) != 2 {
|
|
t.Fatalf("Search(\"red fox\") returned %d hits, want 2", len(hits))
|
|
}
|
|
// Doc 1 matches both terms (higher), doc 3 only 'red'.
|
|
if hits[0].ID != 1 || hits[1].ID != 3 {
|
|
t.Fatalf("hits = %v, want docs [1 3]", hits)
|
|
}
|
|
}
|
|
|
|
func TestIndex_SearchEmptyQuery(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ix := NewIndex[int]()
|
|
ix.Add("hello world", 1)
|
|
|
|
if got := ix.Search(""); len(got) != 0 {
|
|
t.Errorf("Search(\"\") returned %d hits, want 0", len(got))
|
|
}
|
|
if got := ix.Search(" "); len(got) != 0 {
|
|
t.Errorf("Search(\" \") returned %d hits, want 0", len(got))
|
|
}
|
|
}
|
|
|
|
func TestIndex_SearchNoMatch(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ix := NewIndex[int]()
|
|
ix.Add("hello world", 1)
|
|
|
|
if got := ix.Search("unrelatedterm"); len(got) != 0 {
|
|
t.Errorf("Search(unrelatedterm) returned %d hits, want 0", len(got))
|
|
}
|
|
}
|
|
|
|
func TestIndex_ReAddReplacesTerms(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ix := NewIndex[int]()
|
|
ix.Add("hello world", 1)
|
|
ix.Add("goodbye world", 1) // re-add same id
|
|
|
|
if got := ix.Search("hello"); len(got) != 0 {
|
|
t.Errorf("Search(\"hello\") after re-add returned %d hits, want 0", len(got))
|
|
}
|
|
if got := ix.Search("world"); len(got) != 1 {
|
|
t.Errorf("Search(\"world\") after re-add returned %d hits, want 1", len(got))
|
|
}
|
|
}
|
|
|
|
func TestIndex_StringIDs(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ix := NewIndex[string]()
|
|
ix.Add("night sky blue", "moon")
|
|
ix.Add("ocean blue water", "sea")
|
|
|
|
hits := ix.Search("blue")
|
|
if len(hits) != 2 {
|
|
t.Fatalf("Search(\"blue\") returned %d hits, want 2", len(hits))
|
|
}
|
|
if hits[0].ID != "moon" || hits[1].ID != "sea" {
|
|
t.Errorf("hits = %v, want [moon sea]", hits)
|
|
}
|
|
}
|