fork of https://github.com/sourcegraph/zoekt
1package main
2
3import (
4 "crypto/sha1"
5 "fmt"
6 "io"
7 "os"
8 "os/exec"
9 "path/filepath"
10 "strings"
11 "testing"
12
13 "github.com/sourcegraph/zoekt"
14 "github.com/sourcegraph/zoekt/index"
15)
16
17func TestHasMultipleShards(t *testing.T) {
18 dir := t.TempDir()
19
20 cases := []struct {
21 file string
22 wantHasMultipleShards bool
23 }{
24 {"large.00000.zoekt", true},
25 {"large.00001.zoekt", true},
26 {"small.00000.zoekt", false},
27 {"compound-foo.00000.zoekt", false},
28 {"else", false},
29 }
30
31 for _, c := range cases {
32 _, err := os.Create(filepath.Join(dir, c.file))
33 if err != nil {
34 t.Fatal(err)
35 }
36 }
37
38 for _, tt := range cases {
39 t.Run(tt.file, func(t *testing.T) {
40 if got := hasMultipleShards(filepath.Join(dir, tt.file)); got != tt.wantHasMultipleShards {
41 t.Fatalf("want %t, got %t", tt.wantHasMultipleShards, got)
42 }
43 })
44 }
45}
46
47func TestDoNotDeleteSingleShards(t *testing.T) {
48 dir := t.TempDir()
49
50 // Create a test shard.
51 opts := index.Options{
52 IndexDir: dir,
53 RepositoryDescription: zoekt.Repository{Name: "test-repo"},
54 }
55 opts.SetDefaults()
56 b, err := index.NewBuilder(opts)
57 if err != nil {
58 t.Fatalf("NewBuilder: %v", err)
59 }
60 if err := b.AddFile("F", []byte(strings.Repeat("abc", 100))); err != nil {
61 t.Fatalf("AddFile: %v", err)
62 }
63 if err := b.Finish(); err != nil {
64 t.Errorf("Finish: %v", err)
65 }
66
67 s := &Server{IndexDir: dir, mergeOpts: mergeOpts{targetSizeBytes: 2000 * 1024 * 1024}}
68 s.merge(helperCallMerge)
69
70 _, err = os.Stat(filepath.Join(dir, "test-repo_v16.00000.zoekt"))
71 if err != nil {
72 t.Fatal(err)
73 }
74}
75
76func helperCallMerge(s ...string) *exec.Cmd {
77 cs := []string{"-test.run=TestCallMerge", "--"}
78 cs = append(cs, s...)
79 env := []string{
80 "GO_TEST_WANT_CALL_MERGE=1",
81 }
82 cmd := exec.Command(os.Args[0], cs...)
83 cmd.Env = append(env, os.Environ()...)
84 return cmd
85}
86
87func TestCallMerge(t *testing.T) {
88 if os.Getenv("GO_TEST_WANT_CALL_MERGE") != "1" {
89 return
90 }
91 defer os.Exit(0)
92
93 args := os.Args
94 for len(args) > 0 {
95 if args[0] == "--" {
96 args = args[1:]
97 break
98 }
99 args = args[1:]
100 }
101
102 // We mock the merge process by deleting the input shards and creating an empty
103 // compound shard with a proper name.
104 h := sha1.New()
105 for _, a := range args {
106 h.Write([]byte(filepath.Base(a)))
107 h.Write([]byte{0})
108 _ = os.Remove(a)
109 }
110
111 compoundShardName := filepath.Join(filepath.Dir(args[1]), fmt.Sprintf("compound-%x_v%d.%05d.zoekt", h.Sum(nil), 17, 0))
112 f, _ := os.Create(compoundShardName)
113 _ = f.Close()
114
115 // Just like zoekt-merge-index, we write the name of the compound shard to
116 // stdout.
117 _, _ = fmt.Fprint(os.Stdout, compoundShardName)
118}
119
120func TestMerge(t *testing.T) {
121 // A fixed set of shards gives us reliable shard sizes which makes it easy to
122 // define a cutoff with targetSizeBytes.
123 m := []string{
124 "../../testdata/shards/repo_v16.00000.zoekt",
125 "../../testdata/shards/repo2_v16.00000.zoekt",
126 "../../testdata/shards/ctagsrepo_v16.00000.zoekt",
127 }
128
129 testCases := []struct {
130 name string
131 targetSizeBytes int64
132 wantCompound int
133 wantSimple int
134 }{
135 {
136 name: "3 shards",
137 targetSizeBytes: 6 * 1024,
138 wantCompound: 1,
139 wantSimple: 0,
140 },
141 {
142 name: "2 shards",
143 targetSizeBytes: 4 * 1024,
144 wantCompound: 1,
145 wantSimple: 1,
146 },
147 {
148 // This is a pathological case where the target size of a compound shard is
149 // smaller than the size of a simple shard. In realistic scenarios,
150 // targetSizeBytes should be 100x or more of a typical shard size.
151 name: "target size too small",
152 targetSizeBytes: 2 * 1024,
153 wantCompound: 0,
154 wantSimple: 3,
155 },
156 {
157 name: "target size too big",
158 targetSizeBytes: 10 * 1024,
159 wantCompound: 0,
160 wantSimple: 3,
161 },
162 {
163 name: "target size 0",
164 targetSizeBytes: 0,
165 wantCompound: 0,
166 wantSimple: 3,
167 },
168 }
169
170 checkCount := func(dir string, pattern string, want int) {
171 have, err := filepath.Glob(filepath.Join(dir, pattern))
172 if err != nil {
173 t.Fatal(err)
174 }
175 if len(have) != want {
176 t.Fatalf("want %d, have %d", want, len(have))
177 }
178 }
179
180 for _, tc := range testCases {
181 t.Run(tc.name, func(t *testing.T) {
182 dir := t.TempDir()
183 _, err := copyTestShards(dir, m)
184 if err != nil {
185 t.Fatal(err)
186 }
187
188 s := &Server{
189 IndexDir: dir,
190 mergeOpts: mergeOpts{targetSizeBytes: tc.targetSizeBytes},
191 }
192
193 s.merge(helperCallMerge)
194
195 checkCount(dir, "compound-*", tc.wantCompound)
196 checkCount(dir, "*_v16.00000.zoekt", tc.wantSimple)
197 })
198 }
199}
200
201func copyTestShards(dstDir string, srcShards []string) ([]string, error) {
202 var tmpShards []string
203 for _, s := range srcShards {
204 dst := filepath.Join(dstDir, filepath.Base(s))
205 tmpShards = append(tmpShards, dst)
206 if err := copyFile(s, dst); err != nil {
207 return nil, err
208 }
209 }
210 return tmpShards, nil
211}
212
213func copyFile(src, dst string) (err error) {
214 s, err := os.Open(src)
215 if err != nil {
216 return err
217 }
218 defer s.Close()
219
220 d, err := os.Create(dst)
221 if err != nil {
222 return err
223 }
224 if _, err := io.Copy(d, s); err != nil {
225 d.Close()
226 return err
227 }
228 return d.Close()
229}