github.com/keybase/client/go@v0.0.0-20240309051027-028f7c731f8b/kbfs/libkbfs/disk_quota_cache_test.go (about)

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