github.com/zorawar87/trillian@v1.2.1/quota/cacheqm/cache_test.go (about)

     1  // Copyright 2017 Google Inc. All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package cacheqm
    16  
    17  import (
    18  	"context"
    19  	"errors"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/golang/mock/gomock"
    24  	"github.com/google/trillian/quota"
    25  	"github.com/google/trillian/testonly/matchers"
    26  	"github.com/kylelemons/godebug/pretty"
    27  )
    28  
    29  const (
    30  	minBatchSize = 20
    31  	maxEntries   = 10
    32  )
    33  
    34  var (
    35  	specs = []quota.Spec{
    36  		{Group: quota.Global, Kind: quota.Read},
    37  		{Group: quota.Global, Kind: quota.Write},
    38  	}
    39  )
    40  
    41  func TestNewCachedManagerErrors(t *testing.T) {
    42  	tests := []struct {
    43  		minBatchSize, maxEntries int
    44  	}{
    45  		{minBatchSize: 0, maxEntries: 10},
    46  		{minBatchSize: -1, maxEntries: 10},
    47  		{minBatchSize: 10, maxEntries: 0},
    48  		{minBatchSize: 10, maxEntries: -1},
    49  	}
    50  	qm := quota.Noop()
    51  	for _, test := range tests {
    52  		if _, err := NewCachedManager(qm, test.minBatchSize, test.maxEntries); err == nil {
    53  			t.Errorf("NewCachedManager(_, %v, %v) returned err = nil, want non-nil", test.minBatchSize, test.maxEntries)
    54  		}
    55  	}
    56  }
    57  
    58  func TestCachedManager_PeekTokens(t *testing.T) {
    59  	ctrl := gomock.NewController(t)
    60  	defer ctrl.Finish()
    61  
    62  	tests := []struct {
    63  		desc       string
    64  		wantTokens map[quota.Spec]int
    65  		wantErr    error
    66  	}{
    67  		{
    68  			desc: "success",
    69  			wantTokens: map[quota.Spec]int{
    70  				{Group: quota.Global, Kind: quota.Read}: 10,
    71  				{Group: quota.Global, Kind: quota.Read}: 11,
    72  			},
    73  		},
    74  		{
    75  			desc:    "error",
    76  			wantErr: errors.New("llama ate all tokens"),
    77  		},
    78  	}
    79  
    80  	ctx := context.Background()
    81  	for _, test := range tests {
    82  		mock := quota.NewMockManager(ctrl)
    83  		mock.EXPECT().PeekTokens(ctx, specs).Return(test.wantTokens, test.wantErr)
    84  
    85  		qm, err := NewCachedManager(mock, minBatchSize, maxEntries)
    86  		if err != nil {
    87  			t.Fatalf("NewCachedManager() returned err = %v", err)
    88  		}
    89  
    90  		tokens, err := qm.PeekTokens(ctx, specs)
    91  		if diff := pretty.Compare(tokens, test.wantTokens); diff != "" {
    92  			t.Errorf("%v: post-PeekTokens() diff (-got +want):\n%v", test.desc, diff)
    93  		}
    94  		if err != test.wantErr {
    95  			t.Errorf("%v: PeekTokens() returned err = %#v, want = %#v", test.desc, err, test.wantErr)
    96  		}
    97  	}
    98  }
    99  
   100  // TestCachedManager_DelegatedMethods tests all delegated methods that have a single error return.
   101  func TestCachedManager_DelegatedMethods(t *testing.T) {
   102  	ctrl := gomock.NewController(t)
   103  	defer ctrl.Finish()
   104  
   105  	ctx := context.Background()
   106  	tokens := 5
   107  	for _, want := range []error{nil, errors.New("llama ate all tokens")} {
   108  		mock := quota.NewMockManager(ctrl)
   109  		qm, err := NewCachedManager(mock, minBatchSize, maxEntries)
   110  		if err != nil {
   111  			t.Fatalf("NewCachedManager() returned err = %v", err)
   112  		}
   113  
   114  		mock.EXPECT().PutTokens(ctx, tokens, specs).Return(want)
   115  		if err := qm.PutTokens(ctx, tokens, specs); err != want {
   116  			t.Errorf("PutTokens() returned err = %#v, want = %#v", err, want)
   117  		}
   118  
   119  		mock.EXPECT().ResetQuota(ctx, specs).Return(want)
   120  		if err := qm.ResetQuota(ctx, specs); err != want {
   121  			t.Errorf("ResetQuota() returned err = %#v, want = %#v", err, want)
   122  		}
   123  	}
   124  }
   125  
   126  func TestCachedManager_GetTokens_CachesTokens(t *testing.T) {
   127  	ctrl := gomock.NewController(t)
   128  	defer ctrl.Finish()
   129  
   130  	ctx := context.Background()
   131  	tokens := 3
   132  	mock := quota.NewMockManager(ctrl)
   133  	mock.EXPECT().GetTokens(ctx, matchers.AtLeast(minBatchSize), specs).Times(2).Return(nil)
   134  
   135  	qm, err := NewCachedManager(mock, minBatchSize, maxEntries)
   136  	if err != nil {
   137  		t.Fatalf("NewCachedManager() returned err = %v", err)
   138  	}
   139  
   140  	// Quota requests happen in tokens+minBatchSize steps, so that minBatchSize tokens get cached
   141  	// after the the request is satisfied.
   142  	// Therefore, the call pattern below is satisfied by just 2 underlying GetTokens() calls.
   143  	calls := []int{tokens, minBatchSize, tokens, minBatchSize / 2, minBatchSize / 2}
   144  	for i, call := range calls {
   145  		if err := qm.GetTokens(ctx, call, specs); err != nil {
   146  			t.Fatalf("GetTokens() returned err = %v (call #%v)", err, i+1)
   147  		}
   148  	}
   149  }
   150  
   151  func TestCachedManager_GetTokens_EvictsCache(t *testing.T) {
   152  	ctrl := gomock.NewController(t)
   153  	defer ctrl.Finish()
   154  
   155  	mock := quota.NewMockManager(ctrl)
   156  	mock.EXPECT().GetTokens(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
   157  
   158  	ctx := context.Background()
   159  	maxEntries := 100
   160  	qm, err := NewCachedManager(mock, minBatchSize, maxEntries)
   161  	if err != nil {
   162  		t.Fatalf("NewCachedManager() returned err = %v", err)
   163  	}
   164  
   165  	originalNow := now
   166  	defer func() { now = originalNow }()
   167  	currentTime := time.Time{}
   168  	now = func() time.Time {
   169  		currentTime = currentTime.Add(1 * time.Second)
   170  		return currentTime
   171  	}
   172  
   173  	// Ensure Global quotas are the oldest, we don't want those to get evicted regardless of age.
   174  	tokens := 5
   175  	if err := qm.GetTokens(ctx, tokens, []quota.Spec{
   176  		{Group: quota.Global, Kind: quota.Read},
   177  		{Group: quota.Global, Kind: quota.Write},
   178  	}); err != nil {
   179  		t.Fatalf("GetTokens() returned err = %v", err)
   180  	}
   181  
   182  	// Fill the cache up to maxEntries
   183  	firstTree := int64(10)
   184  	tree := firstTree
   185  	for i := 0; i < maxEntries-2; i++ {
   186  		if err := qm.GetTokens(ctx, tokens, treeSpecs(tree)); err != nil {
   187  			t.Fatalf("GetTokens() returned err = %v (i = %v)", err, i)
   188  		}
   189  		tree++
   190  	}
   191  
   192  	// All entries added from now on must cause eviction of the oldest entries.
   193  	// Evict trees in pairs to exercise the inner evict loop.
   194  	evicts := 20
   195  	for i := 0; i < evicts; i += 2 {
   196  		mock.EXPECT().PutTokens(ctx, minBatchSize, treeSpecs(firstTree+int64(i))).Return(nil)
   197  		mock.EXPECT().PutTokens(ctx, minBatchSize, treeSpecs(firstTree+int64(i+1))).Return(nil)
   198  
   199  		specs := []quota.Spec{treeSpec(tree), treeSpec(tree + 1)}
   200  		tree += 2
   201  		if err := qm.GetTokens(ctx, tokens, specs); err != nil {
   202  			t.Fatalf("GetTokens() returned err = %v (i = %v)", err, i)
   203  		}
   204  	}
   205  
   206  	waitChan := make(chan bool, 1)
   207  	go func() {
   208  		qm.(*manager).wait()
   209  		waitChan <- true
   210  	}()
   211  
   212  	select {
   213  	case <-waitChan:
   214  		// OK, test exited cleanly
   215  	case <-time.After(5 * time.Second):
   216  		t.Errorf("Timed out waiting for qm.wait(), failing test")
   217  	}
   218  }
   219  
   220  func TestManager_GetTokensErrors(t *testing.T) {
   221  	ctrl := gomock.NewController(t)
   222  	defer ctrl.Finish()
   223  
   224  	ctx := context.Background()
   225  	want := errors.New("llama ate all tokens")
   226  	mock := quota.NewMockManager(ctrl)
   227  	mock.EXPECT().GetTokens(ctx, gomock.Any(), specs).Return(want)
   228  
   229  	qm, err := NewCachedManager(mock, minBatchSize, maxEntries)
   230  	if err != nil {
   231  		t.Fatalf("NewCachedManager() returned err = %v", err)
   232  	}
   233  
   234  	if err := qm.GetTokens(ctx, 5 /* numTokens */, specs); err != want {
   235  		t.Errorf("GetTokens() returned err = %#v, want = %#v", err, want)
   236  	}
   237  }
   238  
   239  func treeSpecs(treeID int64) []quota.Spec {
   240  	return []quota.Spec{treeSpec(treeID)}
   241  }
   242  
   243  func treeSpec(treeID int64) quota.Spec {
   244  	return quota.Spec{Group: quota.Tree, Kind: quota.Write, TreeID: treeID}
   245  }