forked from philippgille/chromem-go
-
Notifications
You must be signed in to change notification settings - Fork 0
/
query_test.go
108 lines (103 loc) · 2.66 KB
/
query_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package chromem
import (
"reflect"
"slices"
"testing"
)
func TestFilterDocs(t *testing.T) {
docs := map[string]*Document{
"1": {
ID: "1",
Metadata: map[string]string{
"language": "en",
},
Embedding: []float32{0.1, 0.2, 0.3},
Content: "hello world",
},
"2": {
ID: "2",
Metadata: map[string]string{
"language": "de",
},
Embedding: []float32{0.2, 0.3, 0.4},
Content: "hallo welt",
},
}
tt := []struct {
name string
where map[string]string
whereDocument map[string]string
want []*Document
}{
{
name: "meta match",
where: map[string]string{"language": "de"},
whereDocument: nil,
want: []*Document{docs["2"]},
},
{
name: "meta no match",
where: map[string]string{"language": "fr"},
whereDocument: nil,
want: nil,
},
{
name: "content contains all",
where: nil,
whereDocument: map[string]string{"$contains": "llo"},
want: []*Document{docs["1"], docs["2"]},
},
{
name: "content contains one",
where: nil,
whereDocument: map[string]string{"$contains": "hallo"},
want: []*Document{docs["2"]},
},
{
name: "content contains none",
where: nil,
whereDocument: map[string]string{"$contains": "bonjour"},
want: nil,
},
{
name: "content not_contains all",
where: nil,
whereDocument: map[string]string{"$not_contains": "bonjour"},
want: []*Document{docs["1"], docs["2"]},
},
{
name: "content not_contains one",
where: nil,
whereDocument: map[string]string{"$not_contains": "hello"},
want: []*Document{docs["2"]},
},
{
name: "meta and content match",
where: map[string]string{"language": "de"},
whereDocument: map[string]string{"$contains": "hallo"},
want: []*Document{docs["2"]},
},
{
name: "meta + contains + not_contains",
where: map[string]string{"language": "de"},
whereDocument: map[string]string{"$contains": "hallo", "$not_contains": "bonjour"},
want: []*Document{docs["2"]},
},
}
for _, tc := range tt {
t.Run(tc.name, func(t *testing.T) {
got := filterDocs(docs, tc.where, tc.whereDocument)
if !reflect.DeepEqual(got, tc.want) {
// If len is 2, the order might be different (function under test
// is concurrent and order is not guaranteed).
if len(got) == 2 && len(tc.want) == 2 {
slices.Reverse(got)
if reflect.DeepEqual(got, tc.want) {
return
}
}
t.Fatalf("got %v; want %v", got, tc.want)
}
})
}
}