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