1
package sqlite3
2
 
3
import (
4
	"context"
5
	"crypto/rand"
6
	"database/sql"
7
	"os"
8
	"sync"
9
	"testing"
10
	"time"
11
 
12
	"go.miragespace.co/specter/spec/chord"
13
	"go.miragespace.co/specter/spec/protocol"
14
 
15
	"github.com/stretchr/testify/require"
16
	"go.uber.org/zap"
17
	"go.uber.org/zap/zaptest"
18
)
19
 
20
func testGetKV(t *testing.T) *SqliteKV {
21
	t.Helper()
22
 
23
	as := require.New(t)
24
	logger := zaptest.NewLogger(t, zaptest.WrapOptions(zap.AddCaller()))
25
 
26
	dir, err := os.MkdirTemp("", "sql")
27
	as.NoError(err)
28
 
29
	t.Cleanup(func() {
30
		os.RemoveAll(dir)
31
	})
32
 
33
	cfg := Config{
34
		Logger:  logger,
35
		HashFn:  chord.Hash,
36
		DataDir: dir,
37
	}
38
 
39
	kv, err := New(cfg)
40
	as.NoError(err)
41
	t.Cleanup(kv.Close)
42
 
43
	return kv
44
}
45
 
46
func TestInitialize(t *testing.T) {
47
	require.NoError(t, Initialize("cache"))
48
}
49
 
50
func TestCloseClosesDBHandles(t *testing.T) {
51
	as := require.New(t)
52
	kv := testGetKV(t)
53
 
54
	kv.Close()
55
 
56
	as.Error(kv.reader.PingContext(context.Background()))
57
	as.Error(kv.writer.PingContext(context.Background()))
58
 
59
	kv.Close()
60
}
61
 
62
func TestPrepareStatementsReuseReaderExportStatements(t *testing.T) {
63
	as := require.New(t)
64
	kv := testGetKV(t)
65
 
66
	as.Same(kv.stmts.simpleGet, kv.stmts.exportSimpleGet)
67
	as.Same(kv.stmts.prefixList, kv.stmts.exportPrefixList)
68
	as.NotSame(kv.stmts.leaseGet, kv.stmts.exportLeaseGet)
69
}
70
 
71
func TestImportRejectsLengthMismatch(t *testing.T) {
72
	as := require.New(t)
73
	kv := testGetKV(t)
74
 
75
	err := kv.Import(context.Background(), [][]byte{[]byte("a")}, nil)
76
	as.Error(err)
77
	as.ErrorContains(err, "length mismatch")
78
}
79
 
80
func TestImportRejectsNilValue(t *testing.T) {
81
	as := require.New(t)
82
	kv := testGetKV(t)
83
 
84
	err := kv.Import(context.Background(), [][]byte{[]byte("a")}, []*protocol.KVTransfer{nil})
85
	as.Error(err)
86
	as.ErrorContains(err, "values[0] is nil")
87
}
88
 
89
func TestRemoveKeysBatched(t *testing.T) {
90
	as := require.New(t)
91
	kv := testGetKV(t)
92
 
93
	keys := make([][]byte, removeKeysBatchSize+25)
94
	for i := range keys {
95
		key := make([]byte, 16)
96
		value := make([]byte, 8)
97
		rand.Read(key)
98
		rand.Read(value)
99
		keys[i] = key
100
		as.NoError(kv.Put(context.Background(), key, value))
101
	}
102
 
103
	as.NoError(kv.RemoveKeys(context.Background(), keys))
104
 
105
	remaining, err := kv.RangeKeys(context.Background(), 0, 0)
106
	as.NoError(err)
107
	as.Empty(remaining)
108
}
109
 
110
func TestAllKeys(t *testing.T) {
111
	as := require.New(t)
112
	kv := testGetKV(t)
113
 
114
	key := make([]byte, 64)
115
	value := make([]byte, 8)
116
 
117
	num := 1000
118
	for range num {
119
		rand.Read(key)
120
		rand.Read(value)
121
		err := kv.Put(context.Background(), key, value)
122
		as.NoError(err)
123
	}
124
 
125
	keys, err := kv.RangeKeys(context.Background(), 0, 0)
126
	as.NoError(err)
127
	as.Len(keys, num)
128
}
129
 
130
func TestOrderedKeys(t *testing.T) {
131
	as := require.New(t)
132
	kv := testGetKV(t)
133
 
134
	key := make([]byte, 64)
135
	value := make([]byte, 8)
136
 
137
	num := 1000
138
	for range num {
139
		rand.Read(key)
140
		rand.Read(value)
141
		err := kv.Put(context.Background(), key, value)
142
		as.NoError(err)
143
	}
144
 
145
	keys, err := kv.RangeKeys(context.Background(), 0, 0)
146
	as.NoError(err)
147
 
148
	var prev uint64 = 0
149
	for _, key := range keys {
150
		id := chord.Hash(key)
151
		as.LessOrEqual(prev, id)
152
		prev = id
153
	}
154
}
155
 
156
func TestLocalOperations(t *testing.T) {
157
	as := require.New(t)
158
	kv := testGetKV(t)
159
 
160
	num := 32
161
	length := 8
162
	keys := make([][]byte, num)
163
	values := make([]*protocol.KVTransfer, num)
164
 
165
	for i := range keys {
166
		keys[i] = make([]byte, length)
167
		values[i] = &protocol.KVTransfer{
168
			SimpleValue:    make([]byte, length),
169
			PrefixChildren: make([][]byte, 0),
170
		}
171
		rand.Read(keys[i])
172
		rand.Read(values[i].SimpleValue)
173
	}
174
 
175
	as.Nil(kv.Import(context.Background(), keys, values))
176
 
177
	ret, err := kv.Export(context.Background(), keys)
178
	as.NoError(err)
179
	as.EqualValues(values, ret)
180
 
181
	err = kv.RemoveKeys(context.Background(), keys)
182
	as.NoError(err)
183
 
184
	ret, err = kv.Export(context.Background(), keys)
185
	as.NoError(err)
186
	as.NotEqualValues(values, ret)
187
}
188
 
189
func TestComplexImportExport(t *testing.T) {
190
	as := require.New(t)
191
	kv := testGetKV(t)
192
 
193
	key := make([]byte, 8)
194
	rand.Read(key)
195
 
196
	plainValue := make([]byte, 16)
197
	rand.Read(plainValue)
198
	child := make([]byte, 32)
199
	rand.Read(child)
200
 
201
	as.NoError(kv.Put(context.Background(), key, plainValue))
202
	as.NoError(kv.PrefixAppend(context.Background(), key, child))
203
	tk, err := kv.Acquire(context.Background(), key, time.Second)
204
	as.NoError(err)
205
 
206
	val, err := kv.Get(context.Background(), key)
207
	as.NoError(err)
208
	as.EqualValues(plainValue, val)
209
 
210
	vals, err := kv.PrefixList(context.Background(), key)
211
	as.NoError(err)
212
	as.Len(vals, 1)
213
	as.EqualValues(child, vals[0])
214
 
215
	keys, err := kv.RangeKeys(context.Background(), 0, 0)
216
	as.NoError(err)
217
	exp, err := kv.Export(context.Background(), keys)
218
	as.NoError(err)
219
 
220
	kv2 := testGetKV(t)
221
	as.NoError(kv2.Import(context.Background(), keys, exp))
222
 
223
	val, err = kv2.Get(context.Background(), key)
224
	as.NoError(err)
225
	as.EqualValues(plainValue, val)
226
 
227
	vals, err = kv2.PrefixList(context.Background(), key)
228
	as.NoError(err)
229
	as.Len(vals, 1)
230
	as.EqualValues(child, vals[0])
231
 
232
	tk2, err := kv2.Renew(context.Background(), key, time.Second, tk)
233
	as.NoError(err)
234
	as.NoError(kv2.Release(context.Background(), key, tk2))
235
}
236
 
237
// The following test was collaborated with GPT 4o-mini
238
func TestConcurrentOps(t *testing.T) {
239
	kv := testGetKV(t)
240
	as := require.New(t)
241
 
242
	keys := make([][]byte, 8)
243
	for i := range keys {
244
		keys[i] = make([]byte, 8)
245
		rand.Read(keys[i])
246
	}
247
 
248
	var wg sync.WaitGroup
249
	const numGoroutines = 50
250
 
251
	testOperations := func(ctx context.Context, key []byte, id int) {
252
		defer wg.Done()
253
 
254
		err := kv.Put(ctx, key, []byte("test_data"))
255
		as.NoError(err)
256
 
257
		prefixValue := "prefix_value_" + string(rune(id))
258
		err = kv.PrefixAppend(ctx, key, []byte(prefixValue))
259
		as.NoError(err)
260
 
261
		token, err := kv.Acquire(ctx, key, time.Second*5)
262
		if err == nil {
263
			token, err = kv.Renew(ctx, key, time.Second*5, token)
264
			as.NoError(err)
265
 
266
			err = kv.Release(ctx, key, token)
267
			as.NoError(err)
268
		}
269
 
270
		_, err = kv.Get(ctx, key)
271
		as.NoError(err)
272
 
273
		_, err = kv.PrefixList(ctx, key)
274
		as.NoError(err)
275
 
276
		err = kv.PrefixAppend(ctx, key, []byte(prefixValue))
277
		as.Error(err)
278
 
279
		err = kv.Delete(ctx, key)
280
		as.NoError(err)
281
 
282
		err = kv.PrefixRemove(ctx, key, []byte(prefixValue))
283
		as.NoError(err)
284
	}
285
 
286
	ctx := context.Background()
287
	for i := range numGoroutines {
288
		for _, key := range keys {
289
			wg.Add(1)
290
			go testOperations(ctx, key, i)
291
		}
292
	}
293
 
294
	wg.Wait()
295
 
296
	for _, key := range keys {
297
		var flags uint8
298
		err := kv.reader.QueryRow("SELECT `flags` FROM `key_trackers` WHERE `key` = ?", key).Scan(&flags)
299
		as.ErrorIs(err, sql.ErrNoRows)
300
	}
301
 
302
}