From 3b5ceaa190ce8dc161e878b4f50d46c6344c0b1f Mon Sep 17 00:00:00 2001
From: Nikos Angelopoulos <nikolaos.angelopoulos@grafana.com>
Date: Thu, 14 Nov 2024 19:10:40 +0100
Subject: [PATCH] address pr comments

---
 CHANGELOG.md |  1 +
 kv/mock.go   | 54 ++++++++++++++++++----------------------------------
 2 files changed, 20 insertions(+), 35 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d90ec93d1..efd0e0a03 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -237,6 +237,7 @@
 * [ENHANCEMENT] Cache: Add `.Advance()` methods to mock cache clients for easier testing of TTLs. #601
 * [ENHANCEMENT] Memberlist: Add concurrency to the transport's WriteTo method. #525
 * [ENHANCEMENT] Memberlist: Notifications can now be processed once per interval specified by `-memberlist.notify-interval` to reduce notify storms in large clusters. #592
+* [ENHANCEMENT] KV: Add MockCountingClient, which wraps the `kv.client` and can be used in order to count calls at specific functions of the interface. #618
 * [BUGFIX] spanlogger: Support multiple tenant IDs. #59
 * [BUGFIX] Memberlist: fixed corrupted packets when sending compound messages with more than 255 messages or messages bigger than 64KB. #85
 * [BUGFIX] Ring: `ring_member_ownership_percent` and `ring_tokens_owned` metrics are not updated on scale down. #109
diff --git a/kv/mock.go b/kv/mock.go
index 89e30c644..99c84e58d 100644
--- a/kv/mock.go
+++ b/kv/mock.go
@@ -2,7 +2,6 @@ package kv
 
 import (
 	"context"
-	"sync"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -45,72 +44,57 @@ func (m mockClient) WatchPrefix(_ context.Context, _ string, _ func(string, inte
 type MockCountingClient struct {
 	client Client
 
-	mtx   sync.Mutex
-	calls map[string]*atomic.Int32
+	ListCalls        *atomic.Uint32
+	GetCalls         *atomic.Uint32
+	DeleteCalls      *atomic.Uint32
+	CASCalls         *atomic.Uint32
+	WatchKeyCalls    *atomic.Uint32
+	WatchPrefixCalls *atomic.Uint32
 }
 
 func NewMockCountingClient(client Client) *MockCountingClient {
 	return &MockCountingClient{
-		client: client,
-		calls:  make(map[string]*atomic.Int32),
+		client:           client,
+		ListCalls:        atomic.NewUint32(0),
+		GetCalls:         atomic.NewUint32(0),
+		DeleteCalls:      atomic.NewUint32(0),
+		CASCalls:         atomic.NewUint32(0),
+		WatchKeyCalls:    atomic.NewUint32(0),
+		WatchPrefixCalls: atomic.NewUint32(0),
 	}
 }
 
 func (mc *MockCountingClient) List(ctx context.Context, prefix string) ([]string, error) {
-	mc.incCall("List")
+	mc.ListCalls.Inc()
 
 	return mc.client.List(ctx, prefix)
 }
 func (mc *MockCountingClient) Get(ctx context.Context, key string) (interface{}, error) {
-	mc.incCall("Get")
+	mc.GetCalls.Inc()
 
 	return mc.client.Get(ctx, key)
 }
 
 func (mc *MockCountingClient) Delete(ctx context.Context, key string) error {
-	mc.incCall("Delete")
+	mc.DeleteCalls.Inc()
 
 	return mc.client.Delete(ctx, key)
 }
 
 func (mc *MockCountingClient) CAS(ctx context.Context, key string, f func(in interface{}) (out interface{}, retry bool, err error)) error {
-	mc.incCall("CAS")
+	mc.CASCalls.Inc()
 
 	return mc.client.CAS(ctx, key, f)
 }
 
 func (mc *MockCountingClient) WatchKey(ctx context.Context, key string, f func(interface{}) bool) {
-	mc.incCall("WatchKey")
+	mc.WatchKeyCalls.Inc()
 
 	mc.client.WatchKey(ctx, key, f)
 }
 
 func (mc *MockCountingClient) WatchPrefix(ctx context.Context, key string, f func(string, interface{}) bool) {
-	mc.incCall("WatchPrefix")
+	mc.WatchPrefixCalls.Inc()
 
 	mc.client.WatchPrefix(ctx, key, f)
 }
-
-func (mc *MockCountingClient) GetCalls(name string) int {
-	mc.mtx.Lock()
-	defer mc.mtx.Unlock()
-	count, exists := mc.calls[name]
-
-	if !exists {
-		return 0
-	}
-
-	return int(count.Load())
-}
-
-func (mc *MockCountingClient) incCall(name string) {
-	mc.mtx.Lock()
-	defer mc.mtx.Unlock()
-
-	if _, exists := mc.calls[name]; !exists {
-		mc.calls[name] = atomic.NewInt32(0)
-	}
-
-	value := mc.calls[name]
-	value.Add(1)
-}