1package digestset
2
3import (
4	"crypto/sha256"
5	_ "crypto/sha512"
6	"encoding/binary"
7	"math/rand"
8	"testing"
9
10	digest "github.com/opencontainers/go-digest"
11)
12
13func assertEqualDigests(t *testing.T, d1, d2 digest.Digest) {
14	if d1 != d2 {
15		t.Fatalf("Digests do not match:\n\tActual: %s\n\tExpected: %s", d1, d2)
16	}
17}
18
19func TestLookup(t *testing.T) {
20	digests := []digest.Digest{
21		"sha256:1234511111111111111111111111111111111111111111111111111111111111",
22		"sha256:1234111111111111111111111111111111111111111111111111111111111111",
23		"sha256:1234611111111111111111111111111111111111111111111111111111111111",
24		"sha256:5432111111111111111111111111111111111111111111111111111111111111",
25		"sha256:6543111111111111111111111111111111111111111111111111111111111111",
26		"sha256:6432111111111111111111111111111111111111111111111111111111111111",
27		"sha256:6542111111111111111111111111111111111111111111111111111111111111",
28		"sha256:6532111111111111111111111111111111111111111111111111111111111111",
29	}
30
31	dset := NewSet()
32	for i := range digests {
33		if err := dset.Add(digests[i]); err != nil {
34			t.Fatal(err)
35		}
36	}
37
38	dgst, err := dset.Lookup("54")
39	if err != nil {
40		t.Fatal(err)
41	}
42	assertEqualDigests(t, dgst, digests[3])
43
44	_, err = dset.Lookup("1234")
45	if err == nil {
46		t.Fatal("Expected ambiguous error looking up: 1234")
47	}
48	if err != ErrDigestAmbiguous {
49		t.Fatal(err)
50	}
51
52	_, err = dset.Lookup("9876")
53	if err == nil {
54		t.Fatal("Expected not found error looking up: 9876")
55	}
56	if err != ErrDigestNotFound {
57		t.Fatal(err)
58	}
59
60	_, err = dset.Lookup("sha256:1234")
61	if err == nil {
62		t.Fatal("Expected ambiguous error looking up: sha256:1234")
63	}
64	if err != ErrDigestAmbiguous {
65		t.Fatal(err)
66	}
67
68	dgst, err = dset.Lookup("sha256:12345")
69	if err != nil {
70		t.Fatal(err)
71	}
72	assertEqualDigests(t, dgst, digests[0])
73
74	dgst, err = dset.Lookup("sha256:12346")
75	if err != nil {
76		t.Fatal(err)
77	}
78	assertEqualDigests(t, dgst, digests[2])
79
80	dgst, err = dset.Lookup("12346")
81	if err != nil {
82		t.Fatal(err)
83	}
84	assertEqualDigests(t, dgst, digests[2])
85
86	dgst, err = dset.Lookup("12345")
87	if err != nil {
88		t.Fatal(err)
89	}
90	assertEqualDigests(t, dgst, digests[0])
91}
92
93func TestAddDuplication(t *testing.T) {
94	digests := []digest.Digest{
95		"sha256:1234111111111111111111111111111111111111111111111111111111111111",
96		"sha256:1234511111111111111111111111111111111111111111111111111111111111",
97		"sha256:1234611111111111111111111111111111111111111111111111111111111111",
98		"sha256:5432111111111111111111111111111111111111111111111111111111111111",
99		"sha256:6543111111111111111111111111111111111111111111111111111111111111",
100		"sha512:65431111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
101		"sha512:65421111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
102		"sha512:65321111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111",
103	}
104
105	dset := NewSet()
106	for i := range digests {
107		if err := dset.Add(digests[i]); err != nil {
108			t.Fatal(err)
109		}
110	}
111
112	if len(dset.entries) != 8 {
113		t.Fatal("Invalid dset size")
114	}
115
116	if err := dset.Add(digest.Digest("sha256:1234511111111111111111111111111111111111111111111111111111111111")); err != nil {
117		t.Fatal(err)
118	}
119
120	if len(dset.entries) != 8 {
121		t.Fatal("Duplicate digest insert should not increase entries size")
122	}
123
124	if err := dset.Add(digest.Digest("sha384:123451111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111")); err != nil {
125		t.Fatal(err)
126	}
127
128	if len(dset.entries) != 9 {
129		t.Fatal("Insert with different algorithm should be allowed")
130	}
131}
132
133func TestRemove(t *testing.T) {
134	digests, err := createDigests(10)
135	if err != nil {
136		t.Fatal(err)
137	}
138
139	dset := NewSet()
140	for i := range digests {
141		if err := dset.Add(digests[i]); err != nil {
142			t.Fatal(err)
143		}
144	}
145
146	dgst, err := dset.Lookup(digests[0].String())
147	if err != nil {
148		t.Fatal(err)
149	}
150	if dgst != digests[0] {
151		t.Fatalf("Unexpected digest value:\n\tExpected: %s\n\tActual: %s", digests[0], dgst)
152	}
153
154	if err := dset.Remove(digests[0]); err != nil {
155		t.Fatal(err)
156	}
157
158	if _, err := dset.Lookup(digests[0].String()); err != ErrDigestNotFound {
159		t.Fatalf("Expected error %v when looking up removed digest, got %v", ErrDigestNotFound, err)
160	}
161}
162
163func TestAll(t *testing.T) {
164	digests, err := createDigests(100)
165	if err != nil {
166		t.Fatal(err)
167	}
168
169	dset := NewSet()
170	for i := range digests {
171		if err := dset.Add(digests[i]); err != nil {
172			t.Fatal(err)
173		}
174	}
175
176	all := map[digest.Digest]struct{}{}
177	for _, dgst := range dset.All() {
178		all[dgst] = struct{}{}
179	}
180
181	if len(all) != len(digests) {
182		t.Fatalf("Unexpected number of unique digests found:\n\tExpected: %d\n\tActual: %d", len(digests), len(all))
183	}
184
185	for i, dgst := range digests {
186		if _, ok := all[dgst]; !ok {
187			t.Fatalf("Missing element at position %d: %s", i, dgst)
188		}
189	}
190
191}
192
193func assertEqualShort(t *testing.T, actual, expected string) {
194	if actual != expected {
195		t.Fatalf("Unexpected short value:\n\tExpected: %s\n\tActual: %s", expected, actual)
196	}
197}
198
199func TestShortCodeTable(t *testing.T) {
200	digests := []digest.Digest{
201		"sha256:1234111111111111111111111111111111111111111111111111111111111111",
202		"sha256:1234511111111111111111111111111111111111111111111111111111111111",
203		"sha256:1234611111111111111111111111111111111111111111111111111111111111",
204		"sha256:5432111111111111111111111111111111111111111111111111111111111111",
205		"sha256:6543111111111111111111111111111111111111111111111111111111111111",
206		"sha256:6432111111111111111111111111111111111111111111111111111111111111",
207		"sha256:6542111111111111111111111111111111111111111111111111111111111111",
208		"sha256:6532111111111111111111111111111111111111111111111111111111111111",
209	}
210
211	dset := NewSet()
212	for i := range digests {
213		if err := dset.Add(digests[i]); err != nil {
214			t.Fatal(err)
215		}
216	}
217
218	dump := ShortCodeTable(dset, 2)
219
220	if len(dump) < len(digests) {
221		t.Fatalf("Error unexpected size: %d, expecting %d", len(dump), len(digests))
222	}
223	assertEqualShort(t, dump[digests[0]], "12341")
224	assertEqualShort(t, dump[digests[1]], "12345")
225	assertEqualShort(t, dump[digests[2]], "12346")
226	assertEqualShort(t, dump[digests[3]], "54")
227	assertEqualShort(t, dump[digests[4]], "6543")
228	assertEqualShort(t, dump[digests[5]], "64")
229	assertEqualShort(t, dump[digests[6]], "6542")
230	assertEqualShort(t, dump[digests[7]], "653")
231}
232
233func createDigests(count int) ([]digest.Digest, error) {
234	r := rand.New(rand.NewSource(25823))
235	digests := make([]digest.Digest, count)
236	for i := range digests {
237		h := sha256.New()
238		if err := binary.Write(h, binary.BigEndian, r.Int63()); err != nil {
239			return nil, err
240		}
241		digests[i] = digest.NewDigest("sha256", h)
242	}
243	return digests, nil
244}
245
246func benchAddNTable(b *testing.B, n int) {
247	digests, err := createDigests(n)
248	if err != nil {
249		b.Fatal(err)
250	}
251	b.ResetTimer()
252	for i := 0; i < b.N; i++ {
253		dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
254		for j := range digests {
255			if err = dset.Add(digests[j]); err != nil {
256				b.Fatal(err)
257			}
258		}
259	}
260}
261
262func benchLookupNTable(b *testing.B, n int, shortLen int) {
263	digests, err := createDigests(n)
264	if err != nil {
265		b.Fatal(err)
266	}
267	dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
268	for i := range digests {
269		if err := dset.Add(digests[i]); err != nil {
270			b.Fatal(err)
271		}
272	}
273	shorts := make([]string, 0, n)
274	for _, short := range ShortCodeTable(dset, shortLen) {
275		shorts = append(shorts, short)
276	}
277
278	b.ResetTimer()
279	for i := 0; i < b.N; i++ {
280		if _, err = dset.Lookup(shorts[i%n]); err != nil {
281			b.Fatal(err)
282		}
283	}
284}
285
286func benchRemoveNTable(b *testing.B, n int) {
287	digests, err := createDigests(n)
288	if err != nil {
289		b.Fatal(err)
290	}
291	b.ResetTimer()
292	for i := 0; i < b.N; i++ {
293		dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
294		b.StopTimer()
295		for j := range digests {
296			if err = dset.Add(digests[j]); err != nil {
297				b.Fatal(err)
298			}
299		}
300		b.StartTimer()
301		for j := range digests {
302			if err = dset.Remove(digests[j]); err != nil {
303				b.Fatal(err)
304			}
305		}
306	}
307}
308
309func benchShortCodeNTable(b *testing.B, n int, shortLen int) {
310	digests, err := createDigests(n)
311	if err != nil {
312		b.Fatal(err)
313	}
314	dset := &Set{entries: digestEntries(make([]*digestEntry, 0, n))}
315	for i := range digests {
316		if err := dset.Add(digests[i]); err != nil {
317			b.Fatal(err)
318		}
319	}
320
321	b.ResetTimer()
322	for i := 0; i < b.N; i++ {
323		ShortCodeTable(dset, shortLen)
324	}
325}
326
327func BenchmarkAdd10(b *testing.B) {
328	benchAddNTable(b, 10)
329}
330
331func BenchmarkAdd100(b *testing.B) {
332	benchAddNTable(b, 100)
333}
334
335func BenchmarkAdd1000(b *testing.B) {
336	benchAddNTable(b, 1000)
337}
338
339func BenchmarkRemove10(b *testing.B) {
340	benchRemoveNTable(b, 10)
341}
342
343func BenchmarkRemove100(b *testing.B) {
344	benchRemoveNTable(b, 100)
345}
346
347func BenchmarkRemove1000(b *testing.B) {
348	benchRemoveNTable(b, 1000)
349}
350
351func BenchmarkLookup10(b *testing.B) {
352	benchLookupNTable(b, 10, 12)
353}
354
355func BenchmarkLookup100(b *testing.B) {
356	benchLookupNTable(b, 100, 12)
357}
358
359func BenchmarkLookup1000(b *testing.B) {
360	benchLookupNTable(b, 1000, 12)
361}
362
363func BenchmarkShortCode10(b *testing.B) {
364	benchShortCodeNTable(b, 10, 12)
365}
366func BenchmarkShortCode100(b *testing.B) {
367	benchShortCodeNTable(b, 100, 12)
368}
369func BenchmarkShortCode1000(b *testing.B) {
370	benchShortCodeNTable(b, 1000, 12)
371}
372