fork of https://github.com/sourcegraph/zoekt
1package stream
2
3import (
4 "bytes"
5 "context"
6 "encoding/gob"
7 "fmt"
8 "net/http"
9 "net/http/httptest"
10 "testing"
11
12 "github.com/google/go-cmp/cmp"
13
14 "github.com/sourcegraph/zoekt"
15 "github.com/sourcegraph/zoekt/internal/mockSearcher"
16 "github.com/sourcegraph/zoekt/query"
17)
18
19func TestStreamSearch(t *testing.T) {
20 q := query.NewAnd(mustParse("hello world|universe"), query.NewRepoSet("foo/bar", "baz/bam"))
21 searcher := &mockSearcher.MockSearcher{
22 WantSearch: q,
23 SearchResult: &zoekt.SearchResult{
24 Files: []zoekt.FileMatch{
25 {FileName: "bin.go"},
26 },
27 },
28 }
29
30 h := &handler{Searcher: adapter{searcher}}
31
32 s := httptest.NewServer(h)
33 defer s.Close()
34
35 cl := NewClient(s.URL, nil)
36
37 c := make(chan *zoekt.SearchResult, 100)
38
39 err := cl.StreamSearch(context.Background(), q, nil, streamerChan(c))
40 if err != nil {
41 t.Fatal(err)
42 }
43 close(c)
44
45 for res := range c {
46 if res.Files == nil {
47 continue
48 }
49 if res.Files[0].FileName != "bin.go" {
50 t.Errorf("got %s, wanted %s", res.Files[0].FileName, "bin.go")
51 }
52 }
53}
54
55func TestStreamSearchJustStats(t *testing.T) {
56 wantStats := zoekt.Stats{
57 Crashes: 1,
58 }
59 q := query.NewAnd(mustParse("hello world|universe"), query.NewRepoSet("foo/bar", "baz/bam"))
60 searcher := &mockSearcher.MockSearcher{
61 WantSearch: q,
62 SearchResult: &zoekt.SearchResult{
63 Files: []zoekt.FileMatch{},
64 Stats: wantStats,
65 },
66 }
67
68 h := &handler{Searcher: adapter{searcher}}
69
70 s := httptest.NewServer(h)
71 defer s.Close()
72
73 cl := NewClient(s.URL, nil)
74
75 c := make(chan *zoekt.SearchResult, 100)
76
77 err := cl.StreamSearch(context.Background(), q, nil, streamerChan(c))
78 if err != nil {
79 t.Fatal(err)
80 }
81 close(c)
82
83 count := 0
84 for res := range c {
85 count += 1
86 if count > 1 {
87 t.Fatal("expected exactly 1 result, got at least 2")
88 }
89 if d := cmp.Diff(wantStats, res.Stats); d != "" {
90 t.Fatalf("zoekt.Stats mismatch (-want +got): %s\n", d)
91 }
92 }
93 if count != 1 {
94 t.Fatal("expected exactly 1 result, got 0")
95 }
96}
97
98func TestEventStreamWriter(t *testing.T) {
99 registerGob()
100 network := new(bytes.Buffer)
101 enc := gob.NewEncoder(network)
102 dec := gob.NewDecoder(network)
103
104 esw := eventStreamWriter{
105 enc: enc,
106 flush: func() {},
107 }
108
109 tests := []struct {
110 event eventType
111 data interface{}
112 }{
113 {
114 eventDone,
115 nil,
116 },
117 {
118 eventMatches,
119 &zoekt.SearchResult{
120 Files: []zoekt.FileMatch{
121 {FileName: "bin.go"},
122 },
123 },
124 },
125 {
126 eventError,
127 "test error",
128 },
129 }
130
131 for _, tt := range tests {
132 t.Run(tt.event.string(), func(t *testing.T) {
133 err := esw.event(tt.event, tt.data)
134 if err != nil {
135 t.Fatal(err)
136 }
137 reply := new(searchReply)
138 err = dec.Decode(reply)
139 if err != nil {
140 t.Fatal(err)
141 }
142 if reply.Event != tt.event {
143 t.Fatalf("got %s, want %s", reply.Event.string(), tt.event.string())
144 }
145 if d := cmp.Diff(tt.data, reply.Data); d != "" {
146 t.Fatalf("mismatch for event type %s (-want +got):\n%s", tt.event.string(), d)
147 }
148 })
149 }
150}
151
152func TestServerError(t *testing.T) {
153 serverError := fmt.Errorf("zoekt server error")
154 h := func(w http.ResponseWriter, r *http.Request) {
155 esw, err := newEventStreamWriter(w)
156 if err != nil {
157 t.Fatal(err)
158 }
159 err = esw.event(eventError, serverError)
160 if err != nil {
161 t.Fatal(err)
162 }
163 }
164 s := httptest.NewServer(http.HandlerFunc(h))
165 cl := NewClient(s.URL, nil)
166 err := cl.StreamSearch(context.Background(), nil, nil, streamerChan(make(chan *zoekt.SearchResult)))
167 if err == nil {
168 t.Fatalf("got nil, want %s", serverError)
169 }
170}
171
172func mustParse(s string) query.Q {
173 q, err := query.Parse(s)
174 if err != nil {
175 panic(err)
176 }
177 return q
178}
179
180type streamerChan chan<- *zoekt.SearchResult
181
182func (c streamerChan) Send(result *zoekt.SearchResult) {
183 c <- result
184}
185
186type adapter struct {
187 zoekt.Searcher
188}
189
190func (a adapter) StreamSearch(ctx context.Context, q query.Q, opts *zoekt.SearchOptions, sender zoekt.Sender) (err error) {
191 sr, err := a.Searcher.Search(ctx, q, opts)
192 if err != nil {
193 return err
194 }
195 sender.Send(sr)
196 return nil
197}
198
199func TestSamplingStream(t *testing.T) {
200 nonZeroStats := zoekt.Stats{
201 ContentBytesLoaded: 10,
202 }
203 filesEvent := &zoekt.SearchResult{
204 Files: make([]zoekt.FileMatch, 10),
205 Stats: nonZeroStats,
206 }
207 fileEvents := func(n int) []*zoekt.SearchResult {
208 res := make([]*zoekt.SearchResult, n)
209 for i := 0; i < n; i++ {
210 res[i] = filesEvent
211 }
212 return res
213 }
214 statsEvent := &zoekt.SearchResult{
215 Stats: nonZeroStats,
216 }
217 statsEvents := func(n int) []*zoekt.SearchResult {
218 res := make([]*zoekt.SearchResult, n)
219 for i := 0; i < n; i++ {
220 res[i] = statsEvent
221 }
222 return res
223 }
224 cases := []struct {
225 events []*zoekt.SearchResult
226 beforeFlushCount int
227 afterFlushCount int
228 }{
229 // These test cases assume that the sampler only forwards
230 // every 100 stats-only event. In case the sampling logic
231 // changes, these tests are not valuable.
232 {nil, 0, 0},
233 {fileEvents(1), 1, 1},
234 {fileEvents(2), 2, 2},
235 {fileEvents(200), 200, 200},
236 {append(fileEvents(1), statsEvents(1)...), 1, 2},
237 {append(fileEvents(1), statsEvents(2)...), 1, 2},
238 {append(fileEvents(1), statsEvents(99)...), 1, 2},
239 {append(fileEvents(1), statsEvents(100)...), 2, 2},
240 {statsEvents(500), 5, 5},
241 {statsEvents(501), 5, 6},
242 }
243
244 for _, tc := range cases {
245 count := 0
246 ss := NewSamplingSender(SenderFunc(func(*zoekt.SearchResult) {
247 count += 1
248 }))
249
250 for _, event := range tc.events {
251 ss.Send(event)
252 }
253 if count != tc.beforeFlushCount {
254 t.Fatalf("expected %d events, got %d", tc.beforeFlushCount, count)
255 }
256 ss.Flush()
257
258 if count != tc.afterFlushCount {
259 t.Fatalf("expected %d events, got %d", tc.afterFlushCount, count)
260 }
261 }
262}