1// Copyright 2018 Keybase Inc. All rights reserved.
2// Use of this source code is governed by a BSD
3// license that can be found in the LICENSE file.
4
5package libkbfs
6
7import (
8	"io/ioutil"
9	"math/rand"
10	"os"
11	"path/filepath"
12	"testing"
13
14	"github.com/keybase/client/go/kbfs/kbfsblock"
15	"github.com/keybase/client/go/protocol/keybase1"
16	"github.com/stretchr/testify/require"
17	"github.com/syndtr/goleveldb/leveldb/storage"
18	"golang.org/x/net/context"
19)
20
21type testDiskQuotaCacheConfig struct {
22	codecGetter
23	logMaker
24}
25
26func newDiskQuotaCacheLocalForTestWithStorage(
27	t *testing.T, s storage.Storage) *DiskQuotaCacheLocal {
28	cache, err := newDiskQuotaCacheLocalFromStorage(&testDiskQuotaCacheConfig{
29		newTestCodecGetter(),
30		newTestLogMaker(t),
31	}, s, modeTest{modeDefault{}})
32	require.NoError(t, err)
33	err = cache.WaitUntilStarted()
34	require.NoError(t, err)
35	return cache
36}
37
38func newDiskQuotaCacheLocalForTest(t *testing.T) (
39	*DiskQuotaCacheLocal, string) {
40	// Use a disk-based level, instead of memory storage, because we
41	// want to simulate a restart and memory storages can't be reused.
42	tempdir, err := ioutil.TempDir(os.TempDir(), "disk_quota_cache")
43	require.NoError(t, err)
44	s, err := storage.OpenFile(filepath.Join(tempdir, "quota"), false)
45	require.NoError(t, err)
46
47	cache := newDiskQuotaCacheLocalForTestWithStorage(t, s)
48	return cache, tempdir
49}
50
51func shutdownDiskQuotaCacheTest(cache DiskQuotaCache, tempdir string) {
52	cache.Shutdown(context.Background())
53	os.RemoveAll(tempdir)
54}
55
56func makeRandomQuotaWithUsageWrite(t *testing.T) kbfsblock.QuotaInfo {
57	qi := kbfsblock.NewQuotaInfo()
58	qi.Total.Bytes[kbfsblock.UsageWrite] = rand.Int63()
59	return *qi
60}
61
62func TestDiskQuotaCacheCommitAndGet(t *testing.T) {
63	t.Parallel()
64	t.Log("Test that basic quota cache Put and Get operations work.")
65	cache, tempdir := newDiskQuotaCacheLocalForTest(t)
66	defer func() {
67		shutdownDiskQuotaCacheTest(cache, tempdir)
68	}()
69
70	ctx := context.Background()
71	id1 := keybase1.MakeTestUID(1).AsUserOrTeam()
72	qi1 := makeRandomQuotaWithUsageWrite(t)
73
74	t.Log("Put a quota into the cache.")
75	_, err := cache.Get(ctx, id1)
76	require.Error(t, err) // not cached yet
77	err = cache.Put(ctx, id1, qi1)
78	require.NoError(t, err)
79	status := cache.Status(ctx)
80	require.Equal(t, uint64(1), status.NumQuotas)
81
82	t.Log("Get a quota from the cache.")
83	getQI1, err := cache.Get(ctx, id1)
84	require.NoError(t, err)
85	checkWrite := func(a, b kbfsblock.QuotaInfo) {
86		require.Equal(
87			t, a.Total.Bytes[kbfsblock.UsageWrite],
88			b.Total.Bytes[kbfsblock.UsageWrite])
89	}
90	checkWrite(qi1, getQI1)
91
92	t.Log("Check the meters.")
93	status = cache.Status(ctx)
94	require.Equal(t, int64(1), status.Hits.Count)
95	require.Equal(t, int64(1), status.Misses.Count)
96	require.Equal(t, int64(1), status.Puts.Count)
97
98	t.Log("A second entry.")
99	id2 := keybase1.MakeTestTeamID(2, false).AsUserOrTeam()
100	qi2 := makeRandomQuotaWithUsageWrite(t)
101	err = cache.Put(ctx, id2, qi2)
102	require.NoError(t, err)
103	getQI2, err := cache.Get(ctx, id2)
104	require.NoError(t, err)
105	checkWrite(qi2, getQI2)
106
107	t.Log("Override the first user.")
108	qi3 := makeRandomQuotaWithUsageWrite(t)
109	err = cache.Put(ctx, id1, qi3)
110	require.NoError(t, err)
111	getQI3, err := cache.Get(ctx, id1)
112	require.NoError(t, err)
113	checkWrite(qi3, getQI3)
114
115	t.Log("Restart the cache and check the stats")
116	cache.Shutdown(ctx)
117	s, err := storage.OpenFile(filepath.Join(tempdir, "quota"), false)
118	require.NoError(t, err)
119	cache = newDiskQuotaCacheLocalForTestWithStorage(t, s)
120	status = cache.Status(ctx)
121	require.Equal(t, uint64(2), status.NumQuotas)
122	getQI3, err = cache.Get(ctx, id1)
123	require.NoError(t, err)
124	checkWrite(qi3, getQI3)
125	getQI2, err = cache.Get(ctx, id2)
126	require.NoError(t, err)
127	checkWrite(qi2, getQI2)
128}
129