From 97ae790f36be737ac91152a606557f30bb244713 Mon Sep 17 00:00:00 2001
From: Paul Rogers <129207811+paul1r@users.noreply.github.com>
Date: Mon, 4 Mar 2024 15:14:30 -0500
Subject: [PATCH] test: Update chunk/cache/mock to have a GetKeys() function 
 (#12122)

---
 pkg/storage/chunk/cache/mock.go           | 12 ++++++++++++
 pkg/storage/chunk/fetcher/fetcher_test.go | 10 +++++-----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/pkg/storage/chunk/cache/mock.go b/pkg/storage/chunk/cache/mock.go
index d8e62c16059dd..55db7f32a5558 100644
--- 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
index d451372eec124..d73974506d4a1 100644
--- 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