Pass custom time in memory storage constructor

wal-g · Sep 4, 2023 · 1b36239 · 1b36239
1 parent bbfd397
commit 1b36239
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 25 deletions.
diff --git a/internal/backup_list_handler_test.go b/internal/backup_list_handler_test.go
@@ -19,17 +19,13 @@ import (
 )
 
 func TestHandleDefaultBackupList(t *testing.T) {
-	rescueNewObjectTime := memory.NewObjectTime
 	curTime := time.Time{}
-	memory.NewObjectTime = func() time.Time {
+	curTimeFunc := func() time.Time {
 		return curTime.UTC()
 	}
-	defer func() {
-		memory.NewObjectTime = rescueNewObjectTime
-	}()
 
 	t.Run("print correct backups in correct order", func(t *testing.T) {
-		folder := memory.NewFolder("", memory.NewStorage())
+		folder := memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc)))
 		curTime = time.Unix(1690000000, 0)
 		_ = folder.PutObject("base_111_backup_stop_sentinel.json", &bytes.Buffer{})
 		curTime = curTime.Add(time.Second)
@@ -79,11 +75,11 @@ func TestHandleDefaultBackupList(t *testing.T) {
 		memStorages := []cache.NamedFolder{
 			{
 				Name:   "storage_1",
-				Folder: memory.NewFolder("", memory.NewStorage()),
+				Folder: memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc))),
 			},
 			{
 				Name:   "storage_2",
-				Folder: memory.NewFolder("", memory.NewStorage()),
+				Folder: memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc))),
 			},
 		}
 		cacheMock.EXPECT().AllAliveStorages().Return(memStorages, nil)
@@ -126,7 +122,7 @@ func TestHandleDefaultBackupList(t *testing.T) {
 	})
 
 	t.Run("handle error with no backups", func(t *testing.T) {
-		folder := memory.NewFolder("", memory.NewStorage())
+		folder := memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc)))
 
 		infoOutput := new(bytes.Buffer)
 		rescueInfoOutput := tracelog.InfoLogger.Writer()

diff --git a/internal/databases/postgres/backup_list_handler_test.go b/internal/databases/postgres/backup_list_handler_test.go
@@ -19,17 +19,13 @@ import (
 )
 
 func TestHandleDetailedBackupList(t *testing.T) {
-	rescueNewObjectTime := memory.NewObjectTime
 	curTime := time.Time{}
-	memory.NewObjectTime = func() time.Time {
+	curTimeFunc := func() time.Time {
 		return curTime.UTC()
 	}
-	defer func() {
-		memory.NewObjectTime = rescueNewObjectTime
-	}()
 
 	t.Run("print correct backup details in correct order", func(t *testing.T) {
-		folder := memory.NewFolder("", memory.NewStorage())
+		folder := memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc)))
 		curTime = time.Unix(1690000000, 0)
 		_ = folder.PutObject("base_111_backup_stop_sentinel.json", &bytes.Buffer{})
 		_ = folder.PutObject("base_111/metadata.json", bytes.NewBufferString("{}"))
@@ -118,11 +114,11 @@ func TestHandleDetailedBackupList(t *testing.T) {
 		memStorages := []cache.NamedFolder{
 			{
 				Name:   "storage_1",
-				Folder: memory.NewFolder("", memory.NewStorage()),
+				Folder: memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc))),
 			},
 			{
 				Name:   "storage_2",
-				Folder: memory.NewFolder("", memory.NewStorage()),
+				Folder: memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc))),
 			},
 		}
 		cacheMock.EXPECT().AllAliveStorages().Return(memStorages, nil)
@@ -194,7 +190,7 @@ func TestHandleDetailedBackupList(t *testing.T) {
 	})
 
 	t.Run("handle error with no backups", func(t *testing.T) {
-		folder := memory.NewFolder("", memory.NewStorage())
+		folder := memory.NewFolder("", memory.NewStorage(memory.WithCustomTime(curTimeFunc)))
 
 		infoOutput := new(bytes.Buffer)
 		rescueInfoOutput := tracelog.InfoLogger.Writer()

diff --git a/pkg/storages/memory/storage.go b/pkg/storages/memory/storage.go
@@ -21,19 +21,28 @@ type TimeStampedData struct {
 	Size      int
 }
 
-var NewObjectTime = time.Now
-
-func TimeStampData(data bytes.Buffer) TimeStampedData {
-	return TimeStampedData{data, CeilTimeUpToMicroseconds(NewObjectTime()), data.Len()}
+func TimeStampData(data bytes.Buffer, timeNow func() time.Time) TimeStampedData {
+	return TimeStampedData{data, CeilTimeUpToMicroseconds(timeNow()), data.Len()}
 }
 
 // Storage is supposed to be used for tests. It doesn't guarantee data safety!
 type Storage struct {
 	underlying *sync.Map
+	timeNow    func() time.Time
 }
 
-func NewStorage() *Storage {
-	return &Storage{&sync.Map{}}
+func NewStorage(opts ...func(*Storage)) *Storage {
+	s := &Storage{underlying: &sync.Map{}, timeNow: time.Now}
+	for _, o := range opts {
+		o(s)
+	}
+	return s
+}
+
+func WithCustomTime(timeNow func() time.Time) func(*Storage) {
+	return func(s *Storage) {
+		s.timeNow = timeNow
+	}
 }
 
 func (storage *Storage) Load(key string) (value TimeStampedData, exists bool) {
@@ -45,7 +54,7 @@ func (storage *Storage) Load(key string) (value TimeStampedData, exists bool) {
 }
 
 func (storage *Storage) Store(key string, value bytes.Buffer) {
-	storage.underlying.Store(key, TimeStampData(value))
+	storage.underlying.Store(key, TimeStampData(value, storage.timeNow))
 }
 
 func (storage *Storage) Delete(key string) {