test(store/snapshot): fix flaky test (#20078)

Co-authored-by: cool-developer <51834436+cool-develope@users.noreply.github.com>

store/snapshots/store_test.go

@@ -262,6 +262,7 @@ func TestStore_Prune(t *testing.T) {
 }
 
 func TestStore_Save(t *testing.T) {
+	t.Parallel()
 	store := setupStore(t)
 	// Saving a snapshot should work
 	snapshot, err := store.Save(4, 1, makeChunks([][]byte{{1}, {2}}))
@@ -318,19 +319,40 @@ 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)
-	}()
-	mtx.Lock()
-	defer mtx.Unlock()
-	_, err = store.Save(7, 2, makeChunks(nil))
-	require.Error(t, err)
-	_, err = store.Save(8, 1, makeChunks(nil))
-	require.NoError(t, err)
-	close(ch)
+	var (
+		wgStart, wgDone sync.WaitGroup
+		mu              sync.Mutex
+		gotErrHeights   []uint64
+	)
+	srcHeights := []uint64{7, 7, 7, 8, 9}
+	wgStart.Add(len(srcHeights))
+	wgDone.Add(len(srcHeights))
+	for _, h := range srcHeights {
+		ch = make(chan io.ReadCloser, 1)
+		ch <- &ReadCloserMock{} // does not block on a buffered channel
+		close(ch)
+		go func(height uint64) {
+			wgStart.Done()
+			wgStart.Wait() // wait for all routines started
+			if _, err = store.Save(height, 1, ch); err != nil {
+				mu.Lock()
+				gotErrHeights = append(gotErrHeights, height)
+				mu.Unlock()
+			}
+			wgDone.Done()
+		}(h)
+	}
+	wgDone.Wait() // wait for all routines completed
+	assert.Equal(t, []uint64{7, 7}, gotErrHeights)
+}
+
+type ReadCloserMock struct {
+}
+
+func (r ReadCloserMock) Read(p []byte) (n int, err error) {
+	return len(p), io.EOF
+}
+
+func (r ReadCloserMock) Close() error {
+	return nil
 }