test: Update chunk/cache/mock to have a GetKeys() function (#12122)

grafana · Mar 4, 2024 · 97ae790 · 97ae790
1 parent 17988d5
commit 97ae790
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/pkg/storage/chunk/cache/mock.go b/pkg/storage/chunk/cache/mock.go
@@ -12,6 +12,7 @@ type MockCache interface {
 	NumKeyUpdates() int
 	GetInternal() map[string][]byte
 	KeysRequested() int
+	GetKeys() []string
 }
 
 type mockCache struct {
@@ -62,6 +63,17 @@ func (m *mockCache) GetInternal() map[string][]byte {
 	return m.cache
 }
 
+func (m *mockCache) GetKeys() []string {
+	m.Lock()
+	defer m.Unlock()
+
+	keys := make([]string, 0, len(m.cache))
+	for key := range m.cache {
+		keys = append(keys, key)
+	}
+	return keys
+}
+
 func (m *mockCache) KeysRequested() int {
 	return m.keysRequested
 }

diff --git a/pkg/storage/chunk/fetcher/fetcher_test.go b/pkg/storage/chunk/fetcher/fetcher_test.go
@@ -200,11 +200,11 @@ func Test(t *testing.T) {
 			chks, err := f.FetchChunks(context.Background(), test.fetch)
 			assert.NoError(t, err)
 			assertChunks(t, test.fetch, chks)
-			l1actual, err := makeChunksFromMap(c1.GetInternal())
+			l1actual, err := makeChunksFromMapKeys(c1.GetKeys())
 			assert.NoError(t, err)
 			assert.Equal(t, test.l1KeysRequested, c1.KeysRequested())
 			assertChunks(t, test.l1End, l1actual)
-			l2actual, err := makeChunksFromMap(c2.GetInternal())
+			l2actual, err := makeChunksFromMapKeys(c2.GetKeys())
 			assert.NoError(t, err)
 			assert.Equal(t, test.l2KeysRequested, c2.KeysRequested())
 			assertChunks(t, test.l2End, l2actual)
@@ -340,9 +340,9 @@ func makeChunks(now time.Time, tpls ...c) []chunk.Chunk {
 	return chks
 }
 
-func makeChunksFromMap(m map[string][]byte) ([]chunk.Chunk, error) {
-	chks := make([]chunk.Chunk, 0, len(m))
-	for k := range m {
+func makeChunksFromMapKeys(keys []string) ([]chunk.Chunk, error) {
+	chks := make([]chunk.Chunk, 0, len(keys))
+	for _, k := range keys {
 		c, err := chunk.ParseExternalKey("fake", k)
 		if err != nil {
 			return nil, err