cache: Add ability to manually advance "now" for mock cache

Add an `.Advance()` method to MockCache and InstrumentedMockCache to
allow the time considered "now" to be moved without needing to actually
sleep. This is useful for testing when items are set with a TTL and you
would like for them to actually expire as they would in a real cache.

Part of grafana/mimir#9386

Signed-off-by: Nick Pillitteri <nick.pillitteri@grafana.com>

cache/mock.go

@@ -19,25 +19,26 @@ var (
 type MockCache struct {
 	mu    sync.Mutex
 	cache map[string]Item
+	now   time.Time
 }
 
 func NewMockCache() *MockCache {
-	c := &MockCache{}
+	c := &MockCache{now: time.Now()}
 	c.Flush()
 	return c
 }
 
 func (m *MockCache) SetAsync(key string, value []byte, ttl time.Duration) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	m.cache[key] = Item{Data: value, ExpiresAt: time.Now().Add(ttl)}
+	m.cache[key] = Item{Data: value, ExpiresAt: m.now.Add(ttl)}
 }
 
 func (m *MockCache) SetMultiAsync(data map[string][]byte, ttl time.Duration) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	exp := time.Now().Add(ttl)
+	exp := m.now.Add(ttl)
 	for key, val := range data {
 		m.cache[key] = Item{Data: val, ExpiresAt: exp}
 	}
@@ -46,19 +47,19 @@ func (m *MockCache) SetMultiAsync(data map[string][]byte, ttl time.Duration) {
 func (m *MockCache) Set(_ context.Context, key string, value []byte, ttl time.Duration) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	m.cache[key] = Item{Data: value, ExpiresAt: time.Now().Add(ttl)}
+	m.cache[key] = Item{Data: value, ExpiresAt: m.now.Add(ttl)}
 	return nil
 }
 
 func (m *MockCache) Add(_ context.Context, key string, value []byte, ttl time.Duration) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if _, ok := m.cache[key]; ok {
+	if i, ok := m.cache[key]; ok && i.ExpiresAt.After(m.now) {
 		return ErrNotStored
 	}
 
-	m.cache[key] = Item{Data: value, ExpiresAt: time.Now().Add(ttl)}
+	m.cache[key] = Item{Data: value, ExpiresAt: m.now.Add(ttl)}
 	return nil
 }
 
@@ -68,7 +69,7 @@ func (m *MockCache) GetMulti(_ context.Context, keys []string, _ ...Option) map[
 
 	found := make(map[string][]byte, len(keys))
 
-	now := time.Now()
+	now := m.now
 	for _, k := range keys {
 		v, ok := m.cache[k]
 		if ok && now.Before(v.ExpiresAt) {
@@ -114,6 +115,13 @@ func (m *MockCache) Flush() {
 	m.cache = map[string]Item{}
 }
 
+func (m *MockCache) Advance(d time.Duration) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.now = m.now.Add(d)
+}
+
 // InstrumentedMockCache is a mocked cache implementation which also tracks the number
 // of times its functions are called.
 type InstrumentedMockCache struct {
@@ -176,6 +184,10 @@ func (m *InstrumentedMockCache) Flush() {
 	m.cache.Flush()
 }
 
+func (m *InstrumentedMockCache) Advance(d time.Duration) {
+	m.cache.Advance(d)
+}
+
 func (m *InstrumentedMockCache) CountStoreCalls() int {
 	return int(m.storeCount.Load())
 }