From f40496b2ffa26e7e5fa0f5c762c03df70e7d2a11 Mon Sep 17 00:00:00 2001
From: Saketh Kalaga <51327242+renormalize@users.noreply.github.com>
Date: Wed, 11 Sep 2024 14:42:37 +0530
Subject: [PATCH] Defined an interface which mock snapstores should implement
 to emulate tagged snapshots' behavior.

---
 pkg/snapstore/gcs_snapstore_test.go | 24 ++++++++++++++++--------
 pkg/snapstore/snapstore_test.go     | 23 ++++++++++++++++++-----
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/pkg/snapstore/gcs_snapstore_test.go b/pkg/snapstore/gcs_snapstore_test.go
index 899698be6..e4b66ee8a 100644
--- a/pkg/snapstore/gcs_snapstore_test.go
+++ b/pkg/snapstore/gcs_snapstore_test.go
@@ -20,16 +20,24 @@ import (
 // mockGCSClient is a mock client to be used in unit tests.
 type mockGCSClient struct {
 	stiface.Client
-	objects        map[string]*[]byte
-	prefix         string
-	objectMetadata map[string]map[string]string
-	objectMutex    sync.Mutex
+	objects     map[string]*[]byte
+	prefix      string
+	objectTags  map[string]map[string]string
+	objectMutex sync.Mutex
 }
 
 func (m *mockGCSClient) Bucket(name string) stiface.BucketHandle {
 	return &mockBucketHandle{bucket: name, client: m}
 }
 
+func (m *mockGCSClient) setTag(taggedSnapshotName string, tagMap map[string]string) {
+	m.objectTags[taggedSnapshotName] = map[string]string{"x-etcd-snapshot-exclude": "true"}
+}
+
+func (m *mockGCSClient) deleteTag(taggedSnapshotName string) {
+	delete(m.objectTags, taggedSnapshotName)
+}
+
 type mockBucketHandle struct {
 	stiface.BucketHandle
 	bucket string
@@ -48,7 +56,7 @@ func (m *mockBucketHandle) Objects(context.Context, *storage.Query) stiface.Obje
 		keys = append(keys, key)
 	}
 	sort.Strings(keys)
-	return &mockObjectIterator{keys: keys, metadata: m.client.objectMetadata}
+	return &mockObjectIterator{keys: keys, tags: m.client.objectTags}
 }
 
 type mockObjectHandle struct {
@@ -83,7 +91,7 @@ func (m *mockObjectHandle) Delete(context.Context) error {
 	defer m.client.objectMutex.Unlock()
 	if _, ok := m.client.objects[m.object]; ok {
 		delete(m.client.objects, m.object)
-		delete(m.client.objectMetadata, m.object)
+		delete(m.client.objectTags, m.object)
 		return nil
 	}
 	return fmt.Errorf("object %s not found", m.object)
@@ -93,7 +101,7 @@ type mockObjectIterator struct {
 	stiface.ObjectIterator
 	currentIndex int
 	keys         []string
-	metadata     map[string]map[string]string
+	tags         map[string]map[string]string
 }
 
 func (m *mockObjectIterator) Next() (*storage.ObjectAttrs, error) {
@@ -101,7 +109,7 @@ func (m *mockObjectIterator) Next() (*storage.ObjectAttrs, error) {
 		name := m.keys[m.currentIndex]
 		obj := &storage.ObjectAttrs{
 			Name:     name,
-			Metadata: m.metadata[name],
+			Metadata: m.tags[name],
 		}
 		m.currentIndex++
 		return obj, nil
diff --git a/pkg/snapstore/snapstore_test.go b/pkg/snapstore/snapstore_test.go
index db197b989..6a9c0f24e 100644
--- a/pkg/snapstore/snapstore_test.go
+++ b/pkg/snapstore/snapstore_test.go
@@ -41,6 +41,14 @@ type testSnapStore struct {
 	objectCountPerSnapshot int
 }
 
+// tagI is the interface that is to be implemented by mock snapstores to set tags on snapshots
+type tagI interface {
+	// Sets all of the tags for a mocked snapshot
+	setTag(string, map[string]string)
+	// Deletes all of the tags of a mocked snapshot
+	deleteTag(string)
+}
+
 var _ = Describe("Save, List, Fetch, Delete from mock snapstore", func() {
 	var (
 		snap1      brtypes.Snapshot
@@ -103,9 +111,9 @@ var _ = Describe("Save, List, Fetch, Delete from mock snapstore", func() {
 		snap5.GenerateSnapshotName()
 
 		gcsClient = &mockGCSClient{
-			objects:        objectMap,
-			prefix:         prefixV2,
-			objectMetadata: make(map[string]map[string]string),
+			objects:    objectMap,
+			prefix:     prefixV2,
+			objectTags: make(map[string]map[string]string),
 		}
 
 		snapstores = map[string]testSnapStore{
@@ -308,11 +316,16 @@ var _ = Describe("Save, List, Fetch, Delete from mock snapstore", func() {
 				Expect(snapList[secondSnapshotIndex].SnapName).To(Equal(snap5.SnapName))
 
 				// List tests with false and true as arguments only implemented with GCS for now
+				var tag tagI
+				switch provider {
+				case "GCS":
+					tag = gcsClient
+				}
 				if provider == "GCS" {
 					// the tagged snapshot should not be returned by the List() call
 					taggedSnapshot := snapList[0]
 					taggedSnapshotName := path.Join(taggedSnapshot.Prefix, taggedSnapshot.SnapDir, taggedSnapshot.SnapName)
-					gcsClient.objectMetadata[taggedSnapshotName] = map[string]string{"x-etcd-snapshot-exclude": "true"}
+					tag.setTag(taggedSnapshotName, map[string]string{"x-etcd-snapshot-exclude": "true"})
 					snapList, err = snapStore.List(false)
 					Expect(err).ShouldNot(HaveOccurred())
 					Expect(snapList.Len()).Should(Equal((numberSnapshotsInObjectMap - 1) * snapStore.objectCountPerSnapshot))
@@ -325,7 +338,7 @@ var _ = Describe("Save, List, Fetch, Delete from mock snapstore", func() {
 					Expect(snapList[0].SnapName).Should(Equal(taggedSnapshot.SnapName))
 
 					// removing the tag will make the snapshot appear in the List call with false
-					delete(gcsClient.objectMetadata, taggedSnapshotName)
+					tag.deleteTag(taggedSnapshotName)
 					snapList, err = snapStore.List(false)
 					Expect(err).ShouldNot(HaveOccurred())
 					Expect(snapList.Len()).Should(Equal(numberSnapshotsInObjectMap * snapStore.objectCountPerSnapshot))