fork of https://github.com/sourcegraph/zoekt
1package main
2
3import (
4 "archive/tar"
5 "archive/zip"
6 "compress/gzip"
7 "context"
8 "errors"
9 "flag"
10 "fmt"
11 "io"
12 "log"
13 "os"
14 "strings"
15 "testing"
16
17 "github.com/sourcegraph/zoekt"
18 "github.com/sourcegraph/zoekt/build"
19 "github.com/sourcegraph/zoekt/query"
20 "github.com/sourcegraph/zoekt/shards"
21)
22
23func TestMain(m *testing.M) {
24 flag.Parse()
25 if !testing.Verbose() {
26 log.SetOutput(io.Discard)
27 }
28 os.Exit(m.Run())
29}
30
31func writeArchive(w io.Writer, format string, files map[string]string) (err error) {
32 if format == "zip" {
33 zw := zip.NewWriter(w)
34 for name, body := range files {
35 f, err := zw.Create(name)
36 if err != nil {
37 return err
38 }
39 if _, err := f.Write([]byte(body)); err != nil {
40 return err
41 }
42 }
43 return zw.Close()
44 }
45
46 if format == "tgz" {
47 gw := gzip.NewWriter(w)
48 defer func() {
49 err2 := gw.Close()
50 if err == nil {
51 err = err2
52 }
53 }()
54 w = gw
55 format = "tar"
56 }
57
58 if format != "tar" {
59 return errors.New("expected tar")
60 }
61
62 tw := tar.NewWriter(w)
63
64 for name, body := range files {
65 hdr := &tar.Header{
66 Name: name,
67 Mode: 0o600,
68 Size: int64(len(body)),
69 }
70 if err := tw.WriteHeader(hdr); err != nil {
71 return err
72 }
73 if _, err := tw.Write([]byte(body)); err != nil {
74 return err
75 }
76 }
77 if err := tw.Close(); err != nil {
78 return err
79 }
80
81 return nil
82}
83
84// TestIndexArg tests zoekt-archive-index by creating an archive and then
85// indexing and executing searches and checking we get expected results.
86// Additionally, we test that the index is properly updated with the
87// -incremental=true option changing the options between indexes and ensuring
88// the results change as expected.
89func TestIndexIncrementally(t *testing.T) {
90 for _, format := range []string{"tar", "tgz", "zip"} {
91 t.Run(format, func(t *testing.T) {
92 testIndexIncrementally(t, format)
93 })
94 }
95}
96
97func testIndexIncrementally(t *testing.T, format string) {
98 indexDir := t.TempDir()
99
100 archive, err := os.CreateTemp("", "TestIndexArg-archive")
101 if err != nil {
102 t.Fatalf("TempFile: %v", err)
103 }
104 defer os.Remove(archive.Name())
105
106 fileSize := 1000
107
108 files := map[string]string{}
109 for i := 0; i < 4; i++ {
110 s := fmt.Sprintf("%d", i)
111 files["F"+s] = strings.Repeat("a", fileSize)
112 files["!F"+s] = strings.Repeat("a", fileSize)
113 }
114
115 err = writeArchive(archive, format, files)
116 if err != nil {
117 t.Fatalf("unable to create archive %v", err)
118 }
119 archive.Close()
120
121 // tests contain options used to build an index and the expected number of
122 // files in the result set based on the options.
123 tests := []struct {
124 largeFiles []string
125 wantNumFiles int
126 }{
127 {
128 largeFiles: []string{},
129 wantNumFiles: 0,
130 },
131 {
132 largeFiles: []string{"F0", "F2"},
133 wantNumFiles: 2,
134 },
135 {
136 largeFiles: []string{"F?", "!F2"},
137 wantNumFiles: 3,
138 },
139 {
140 largeFiles: []string{"F?", "!F2", "\\!F0"},
141 wantNumFiles: 4,
142 },
143 {
144 largeFiles: []string{"F?", "!F2", "\\!F0", "F2"},
145 wantNumFiles: 5,
146 },
147 }
148
149 for _, test := range tests {
150 largeFiles, wantNumFiles := test.largeFiles, test.wantNumFiles
151
152 bopts := build.Options{
153 SizeMax: fileSize - 1,
154 IndexDir: indexDir,
155 LargeFiles: largeFiles,
156 }
157 opts := Options{
158 Incremental: true,
159 Archive: archive.Name(),
160 Name: "repo",
161 Branch: "master",
162 Commit: "cccccccccccccccccccccccccccccccccccccccc",
163 Strip: 0,
164 }
165
166 if err := do(opts, bopts); err != nil {
167 t.Fatalf("error creating index: %v", err)
168 }
169
170 ss, err := shards.NewDirectorySearcher(indexDir)
171 if err != nil {
172 t.Fatalf("NewDirectorySearcher(%s): %v", indexDir, err)
173 }
174 defer ss.Close()
175
176 q, err := query.Parse("aaa")
177 if err != nil {
178 t.Fatalf("Parse(aaa): %v", err)
179 }
180
181 var sOpts zoekt.SearchOptions
182 result, err := ss.Search(context.Background(), q, &sOpts)
183 if err != nil {
184 t.Fatalf("Search(%v): %v", q, err)
185 }
186
187 if len(result.Files) != wantNumFiles {
188 t.Errorf("got %v, want %d files.", result.Files, wantNumFiles)
189 }
190 }
191}