From 25af567521889893a26836e0a3a4e0390b6da6d2 Mon Sep 17 00:00:00 2001
From: cool-developer <51834436+cool-develope@users.noreply.github.com>
Date: Mon, 18 Mar 2024 12:23:22 -0400
Subject: [PATCH] chore: fix the flaky test in store/snapshots (#19774)

Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>
---
 store/snapshots/helpers_test.go | 14 --------------
 store/snapshots/manager_test.go |  2 +-
 store/snapshots/store_test.go   | 12 ++++++++----
 3 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/store/snapshots/helpers_test.go b/store/snapshots/helpers_test.go
index a0d1add9a8..f01a2b6352 100644
--- a/store/snapshots/helpers_test.go
+++ b/store/snapshots/helpers_test.go
@@ -7,7 +7,6 @@ import (
 	"crypto/sha256"
 	"errors"
 	"io"
-	"os"
 	"testing"
 	"time"
 
@@ -288,16 +287,3 @@ func (s *extSnapshotter) RestoreExtension(height uint64, format uint32, payloadR
 	// finalize restoration
 	return nil
 }
-
-// GetTempDir returns a writable temporary director for the test to use.
-func GetTempDir(tb testing.TB) string {
-	tb.Helper()
-	// os.MkDir() is used instead of testing.T.TempDir()
-	// see https://github.com/cosmos/cosmos-sdk/pull/8475 and
-	// https://github.com/cosmos/cosmos-sdk/pull/10341 for
-	// this change's rationale.
-	tempdir, err := os.MkdirTemp("", "")
-	require.NoError(tb, err)
-	tb.Cleanup(func() { _ = os.RemoveAll(tempdir) })
-	return tempdir
-}
diff --git a/store/snapshots/manager_test.go b/store/snapshots/manager_test.go
index 987d0c3b0e..b06c282163 100644
--- a/store/snapshots/manager_test.go
+++ b/store/snapshots/manager_test.go
@@ -236,7 +236,7 @@ func TestManager_Restore(t *testing.T) {
 
 func TestManager_TakeError(t *testing.T) {
 	snapshotter := &mockErrorCommitSnapshotter{}
-	store, err := snapshots.NewStore(GetTempDir(t))
+	store, err := snapshots.NewStore(t.TempDir())
 	require.NoError(t, err)
 	manager := snapshots.NewManager(store, opts, snapshotter, &mockStorageSnapshotter{}, nil, log.NewNopLogger())
 
diff --git a/store/snapshots/store_test.go b/store/snapshots/store_test.go
index c6708ec8d7..45fbf921c9 100644
--- a/store/snapshots/store_test.go
+++ b/store/snapshots/store_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"io"
+	"sync"
 	"testing"
 	"time"
 
@@ -16,7 +17,7 @@ import (
 
 func setupStore(t *testing.T) *snapshots.Store {
 	t.Helper()
-	store, err := snapshots.NewStore(GetTempDir(t))
+	store, err := snapshots.NewStore(t.TempDir())
 	require.NoError(t, err)
 
 	_, err = store.Save(1, 1, makeChunks([][]byte{
@@ -40,8 +41,7 @@ func setupStore(t *testing.T) *snapshots.Store {
 }
 
 func TestNewStore(t *testing.T) {
-	tempdir := GetTempDir(t)
-	_, err := snapshots.NewStore(tempdir)
+	_, err := snapshots.NewStore(t.TempDir())
 
 	require.NoError(t, err)
 }
@@ -319,11 +319,15 @@ func TestStore_Save(t *testing.T) {
 	// Saving a snapshot should error if a snapshot is already in progress for the same height,
 	// regardless of format. However, a different height should succeed.
 	ch = make(chan io.ReadCloser)
+	mtx := sync.Mutex{}
+	mtx.Lock()
 	go func() {
+		mtx.Unlock()
 		_, err := store.Save(7, 1, ch)
 		require.NoError(t, err)
 	}()
-	time.Sleep(10 * time.Millisecond)
+	mtx.Lock()
+	defer mtx.Unlock()
 	_, err = store.Save(7, 2, makeChunks(nil))
 	require.Error(t, err)
 	_, err = store.Save(8, 1, makeChunks(nil))