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