1package main_test
2
3import (
4	"bytes"
5	crypto "crypto/rand"
6	"encoding/binary"
7	"fmt"
8	"io"
9	"io/ioutil"
10	"math/rand"
11	"os"
12	"strconv"
13	"testing"
14
15	"github.com/boltdb/bolt"
16	"github.com/boltdb/bolt/cmd/bolt"
17)
18
19// Ensure the "info" command can print information about a database.
20func TestInfoCommand_Run(t *testing.T) {
21	db := MustOpen(0666, nil)
22	db.DB.Close()
23	defer db.Close()
24
25	// Run the info command.
26	m := NewMain()
27	if err := m.Run("info", db.Path); err != nil {
28		t.Fatal(err)
29	}
30}
31
32// Ensure the "stats" command executes correctly with an empty database.
33func TestStatsCommand_Run_EmptyDatabase(t *testing.T) {
34	// Ignore
35	if os.Getpagesize() != 4096 {
36		t.Skip("system does not use 4KB page size")
37	}
38
39	db := MustOpen(0666, nil)
40	defer db.Close()
41	db.DB.Close()
42
43	// Generate expected result.
44	exp := "Aggregate statistics for 0 buckets\n\n" +
45		"Page count statistics\n" +
46		"\tNumber of logical branch pages: 0\n" +
47		"\tNumber of physical branch overflow pages: 0\n" +
48		"\tNumber of logical leaf pages: 0\n" +
49		"\tNumber of physical leaf overflow pages: 0\n" +
50		"Tree statistics\n" +
51		"\tNumber of keys/value pairs: 0\n" +
52		"\tNumber of levels in B+tree: 0\n" +
53		"Page size utilization\n" +
54		"\tBytes allocated for physical branch pages: 0\n" +
55		"\tBytes actually used for branch data: 0 (0%)\n" +
56		"\tBytes allocated for physical leaf pages: 0\n" +
57		"\tBytes actually used for leaf data: 0 (0%)\n" +
58		"Bucket statistics\n" +
59		"\tTotal number of buckets: 0\n" +
60		"\tTotal number on inlined buckets: 0 (0%)\n" +
61		"\tBytes used for inlined buckets: 0 (0%)\n"
62
63	// Run the command.
64	m := NewMain()
65	if err := m.Run("stats", db.Path); err != nil {
66		t.Fatal(err)
67	} else if m.Stdout.String() != exp {
68		t.Fatalf("unexpected stdout:\n\n%s", m.Stdout.String())
69	}
70}
71
72// Ensure the "stats" command can execute correctly.
73func TestStatsCommand_Run(t *testing.T) {
74	// Ignore
75	if os.Getpagesize() != 4096 {
76		t.Skip("system does not use 4KB page size")
77	}
78
79	db := MustOpen(0666, nil)
80	defer db.Close()
81
82	if err := db.Update(func(tx *bolt.Tx) error {
83		// Create "foo" bucket.
84		b, err := tx.CreateBucket([]byte("foo"))
85		if err != nil {
86			return err
87		}
88		for i := 0; i < 10; i++ {
89			if err := b.Put([]byte(strconv.Itoa(i)), []byte(strconv.Itoa(i))); err != nil {
90				return err
91			}
92		}
93
94		// Create "bar" bucket.
95		b, err = tx.CreateBucket([]byte("bar"))
96		if err != nil {
97			return err
98		}
99		for i := 0; i < 100; i++ {
100			if err := b.Put([]byte(strconv.Itoa(i)), []byte(strconv.Itoa(i))); err != nil {
101				return err
102			}
103		}
104
105		// Create "baz" bucket.
106		b, err = tx.CreateBucket([]byte("baz"))
107		if err != nil {
108			return err
109		}
110		if err := b.Put([]byte("key"), []byte("value")); err != nil {
111			return err
112		}
113
114		return nil
115	}); err != nil {
116		t.Fatal(err)
117	}
118	db.DB.Close()
119
120	// Generate expected result.
121	exp := "Aggregate statistics for 3 buckets\n\n" +
122		"Page count statistics\n" +
123		"\tNumber of logical branch pages: 0\n" +
124		"\tNumber of physical branch overflow pages: 0\n" +
125		"\tNumber of logical leaf pages: 1\n" +
126		"\tNumber of physical leaf overflow pages: 0\n" +
127		"Tree statistics\n" +
128		"\tNumber of keys/value pairs: 111\n" +
129		"\tNumber of levels in B+tree: 1\n" +
130		"Page size utilization\n" +
131		"\tBytes allocated for physical branch pages: 0\n" +
132		"\tBytes actually used for branch data: 0 (0%)\n" +
133		"\tBytes allocated for physical leaf pages: 4096\n" +
134		"\tBytes actually used for leaf data: 1996 (48%)\n" +
135		"Bucket statistics\n" +
136		"\tTotal number of buckets: 3\n" +
137		"\tTotal number on inlined buckets: 2 (66%)\n" +
138		"\tBytes used for inlined buckets: 236 (11%)\n"
139
140	// Run the command.
141	m := NewMain()
142	if err := m.Run("stats", db.Path); err != nil {
143		t.Fatal(err)
144	} else if m.Stdout.String() != exp {
145		t.Fatalf("unexpected stdout:\n\n%s", m.Stdout.String())
146	}
147}
148
149// Main represents a test wrapper for main.Main that records output.
150type Main struct {
151	*main.Main
152	Stdin  bytes.Buffer
153	Stdout bytes.Buffer
154	Stderr bytes.Buffer
155}
156
157// NewMain returns a new instance of Main.
158func NewMain() *Main {
159	m := &Main{Main: main.NewMain()}
160	m.Main.Stdin = &m.Stdin
161	m.Main.Stdout = &m.Stdout
162	m.Main.Stderr = &m.Stderr
163	return m
164}
165
166// MustOpen creates a Bolt database in a temporary location.
167func MustOpen(mode os.FileMode, options *bolt.Options) *DB {
168	// Create temporary path.
169	f, _ := ioutil.TempFile("", "bolt-")
170	f.Close()
171	os.Remove(f.Name())
172
173	db, err := bolt.Open(f.Name(), mode, options)
174	if err != nil {
175		panic(err.Error())
176	}
177	return &DB{DB: db, Path: f.Name()}
178}
179
180// DB is a test wrapper for bolt.DB.
181type DB struct {
182	*bolt.DB
183	Path string
184}
185
186// Close closes and removes the database.
187func (db *DB) Close() error {
188	defer os.Remove(db.Path)
189	return db.DB.Close()
190}
191
192func TestCompactCommand_Run(t *testing.T) {
193	var s int64
194	if err := binary.Read(crypto.Reader, binary.BigEndian, &s); err != nil {
195		t.Fatal(err)
196	}
197	rand.Seed(s)
198
199	dstdb := MustOpen(0666, nil)
200	dstdb.Close()
201
202	// fill the db
203	db := MustOpen(0666, nil)
204	if err := db.Update(func(tx *bolt.Tx) error {
205		n := 2 + rand.Intn(5)
206		for i := 0; i < n; i++ {
207			k := []byte(fmt.Sprintf("b%d", i))
208			b, err := tx.CreateBucketIfNotExists(k)
209			if err != nil {
210				return err
211			}
212			if err := b.SetSequence(uint64(i)); err != nil {
213				return err
214			}
215			if err := fillBucket(b, append(k, '.')); err != nil {
216				return err
217			}
218		}
219		return nil
220	}); err != nil {
221		db.Close()
222		t.Fatal(err)
223	}
224
225	// make the db grow by adding large values, and delete them.
226	if err := db.Update(func(tx *bolt.Tx) error {
227		b, err := tx.CreateBucketIfNotExists([]byte("large_vals"))
228		if err != nil {
229			return err
230		}
231		n := 5 + rand.Intn(5)
232		for i := 0; i < n; i++ {
233			v := make([]byte, 1000*1000*(1+rand.Intn(5)))
234			_, err := crypto.Read(v)
235			if err != nil {
236				return err
237			}
238			if err := b.Put([]byte(fmt.Sprintf("l%d", i)), v); err != nil {
239				return err
240			}
241		}
242		return nil
243	}); err != nil {
244		db.Close()
245		t.Fatal(err)
246	}
247	if err := db.Update(func(tx *bolt.Tx) error {
248		c := tx.Bucket([]byte("large_vals")).Cursor()
249		for k, _ := c.First(); k != nil; k, _ = c.Next() {
250			if err := c.Delete(); err != nil {
251				return err
252			}
253		}
254		return tx.DeleteBucket([]byte("large_vals"))
255	}); err != nil {
256		db.Close()
257		t.Fatal(err)
258	}
259	db.DB.Close()
260	defer db.Close()
261	defer dstdb.Close()
262
263	dbChk, err := chkdb(db.Path)
264	if err != nil {
265		t.Fatal(err)
266	}
267
268	m := NewMain()
269	if err := m.Run("compact", "-o", dstdb.Path, db.Path); err != nil {
270		t.Fatal(err)
271	}
272
273	dbChkAfterCompact, err := chkdb(db.Path)
274	if err != nil {
275		t.Fatal(err)
276	}
277
278	dstdbChk, err := chkdb(dstdb.Path)
279	if err != nil {
280		t.Fatal(err)
281	}
282
283	if !bytes.Equal(dbChk, dbChkAfterCompact) {
284		t.Error("the original db has been touched")
285	}
286	if !bytes.Equal(dbChk, dstdbChk) {
287		t.Error("the compacted db data isn't the same than the original db")
288	}
289}
290
291func fillBucket(b *bolt.Bucket, prefix []byte) error {
292	n := 10 + rand.Intn(50)
293	for i := 0; i < n; i++ {
294		v := make([]byte, 10*(1+rand.Intn(4)))
295		_, err := crypto.Read(v)
296		if err != nil {
297			return err
298		}
299		k := append(prefix, []byte(fmt.Sprintf("k%d", i))...)
300		if err := b.Put(k, v); err != nil {
301			return err
302		}
303	}
304	// limit depth of subbuckets
305	s := 2 + rand.Intn(4)
306	if len(prefix) > (2*s + 1) {
307		return nil
308	}
309	n = 1 + rand.Intn(3)
310	for i := 0; i < n; i++ {
311		k := append(prefix, []byte(fmt.Sprintf("b%d", i))...)
312		sb, err := b.CreateBucket(k)
313		if err != nil {
314			return err
315		}
316		if err := fillBucket(sb, append(k, '.')); err != nil {
317			return err
318		}
319	}
320	return nil
321}
322
323func chkdb(path string) ([]byte, error) {
324	db, err := bolt.Open(path, 0666, nil)
325	if err != nil {
326		return nil, err
327	}
328	defer db.Close()
329	var buf bytes.Buffer
330	err = db.View(func(tx *bolt.Tx) error {
331		return tx.ForEach(func(name []byte, b *bolt.Bucket) error {
332			return walkBucket(b, name, nil, &buf)
333		})
334	})
335	if err != nil {
336		return nil, err
337	}
338	return buf.Bytes(), nil
339}
340
341func walkBucket(parent *bolt.Bucket, k []byte, v []byte, w io.Writer) error {
342	if _, err := fmt.Fprintf(w, "%d:%x=%x\n", parent.Sequence(), k, v); err != nil {
343		return err
344	}
345
346	// not a bucket, exit.
347	if v != nil {
348		return nil
349	}
350	return parent.ForEach(func(k, v []byte) error {
351		if v == nil {
352			return walkBucket(parent.Bucket(k), k, nil, w)
353		}
354		return walkBucket(parent, k, v, w)
355	})
356}
357