1package main
2
3import (
4 "context"
5 "encoding/csv"
6 "log/slog"
7 "strings"
8 "testing"
9
10 "github.com/cockroachdb/pebble"
11 "github.com/cockroachdb/pebble/vfs"
12 "github.com/stretchr/testify/assert"
13)
14
15type debugWriter struct {
16 t *testing.T
17}
18
19func (w *debugWriter) Write(p []byte) (n int, err error) {
20 w.t.Helper()
21 w.t.Log(string(p))
22 return len(p), nil
23}
24
25// make a new pebble that writes to memory and logs to test.Log
26func newMem(t *testing.T) *PebbleCollectionDirectory {
27 memfs := vfs.NewMem()
28 db, err := pebble.Open("wat", &pebble.Options{
29 FS: memfs,
30 })
31 if err != nil {
32 panic(err)
33 }
34
35 log := slog.New(slog.NewTextHandler(&debugWriter{t: t}, &slog.HandlerOptions{Level: slog.LevelDebug}))
36 pcd := &PebbleCollectionDirectory{
37 db: db,
38 collections: make(map[string]uint32),
39 collectionNames: make(map[uint32]string),
40 log: log,
41 }
42 if pcd.log == nil {
43 pcd.log = slog.Default()
44 }
45 return pcd
46}
47
48// did, collection
49const testDataCsv = `alice,post
50alice,like
51bob,post
52bob,other
53carol,post
54eve,post
55eve,like
56eve,other`
57
58func TestPebbleCollectionDirectory(t *testing.T) {
59 assert := assert.New(t)
60
61 pcd := newMem(t)
62 defer func() {
63 err := pcd.Close()
64 if err != nil {
65 t.Error(err)
66 }
67 }()
68
69 rows, err := csv.NewReader(strings.NewReader(testDataCsv)).ReadAll()
70 assert.NoError(err)
71 for _, row := range rows {
72 err := pcd.MaybeSetCollection(row[0], row[1])
73 assert.NoError(err)
74 }
75 stats, err := pcd.GetCollectionStats()
76 assert.NoError(err)
77 t.Log(stats)
78 assert.Equal(uint64(4), stats.CollectionCounts["post"])
79 assert.Equal(uint64(2), stats.CollectionCounts["like"])
80 assert.Equal(uint64(2), stats.CollectionCounts["other"])
81
82 t.Log(pcd.collections)
83
84 wat, nextCursor, err := pcd.ReadCollection(context.Background(), "post", "", 1000)
85 assert.NoError(err)
86 assert.Equal("", nextCursor)
87 for _, row := range wat {
88 assert.Equal("post", row.Collection)
89 }
90 assert.Equal(4, len(wat))
91
92 wat, nextCursor, err = pcd.ReadCollection(context.Background(), "like", "", 1000)
93 assert.NoError(err)
94 assert.Equal("", nextCursor)
95 for _, row := range wat {
96 assert.Equal("like", row.Collection)
97 }
98 assert.Equal(2, len(wat))
99
100 wat, nextCursor, err = pcd.ReadCollection(context.Background(), "other", "", 1000)
101 assert.NoError(err)
102 assert.Equal("", nextCursor)
103 for _, row := range wat {
104 assert.Equal("other", row.Collection)
105 }
106 assert.Equal(2, len(wat))
107}