From 4f9167967f4ac28c0ce82faff81ebf572f4b9862 Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Wed, 13 Nov 2024 10:25:11 +0800
Subject: [PATCH] Problem: async prune test is failed

---
 store/rootmulti/store_test.go | 49 ++++++++++++++++++++++-------------
 1 file changed, 31 insertions(+), 18 deletions(-)

diff --git a/store/rootmulti/store_test.go b/store/rootmulti/store_test.go
index 3d7dca809d2f..dcb5d8482af7 100644
--- a/store/rootmulti/store_test.go
+++ b/store/rootmulti/store_test.go
@@ -532,14 +532,18 @@ func TestMultiStore_Pruning(t *testing.T) {
 				ms.Commit()
 			}
 
-			for _, v := range tc.saved {
-				_, err := ms.CacheMultiStoreWithVersion(v)
-				require.NoError(t, err, "expected no error when loading height: %d", v)
+			for _, v := range tc.deleted {
+				// Ensure async pruning is done
+				checkErr := func() bool {
+					_, err := ms.CacheMultiStoreWithVersion(v)
+					return err != nil
+				}
+				require.Eventually(t, checkErr, 1*time.Second, 10*time.Millisecond, "expected error when loading height: %d", v)
 			}
 
-			for _, v := range tc.deleted {
+			for _, v := range tc.saved {
 				_, err := ms.CacheMultiStoreWithVersion(v)
-				require.Error(t, err, "expected error when loading height: %d", v)
+				require.NoError(t, err, "expected no error when loading height: %d", v)
 			}
 		})
 	}
@@ -564,16 +568,6 @@ func TestMultiStore_Pruning_SameHeightsTwice(t *testing.T) {
 
 	require.Equal(t, numVersions, lastCommitInfo.Version)
 
-	for v := int64(1); v < numVersions-int64(keepRecent); v++ {
-		err := ms.LoadVersion(v)
-		require.Error(t, err, "expected error when loading pruned height: %d", v)
-	}
-
-	for v := (numVersions - int64(keepRecent)); v < numVersions; v++ {
-		err := ms.LoadVersion(v)
-		require.NoError(t, err, "expected no error when loading height: %d", v)
-	}
-
 	// Get latest
 	err := ms.LoadVersion(numVersions - 1)
 	require.NoError(t, err)
@@ -588,6 +582,17 @@ func TestMultiStore_Pruning_SameHeightsTwice(t *testing.T) {
 	// Ensure that can commit one more height with no panic
 	lastCommitInfo = ms.Commit()
 	require.Equal(t, numVersions+1, lastCommitInfo.Version)
+	isPruned := func() bool {
+		ls := ms.Commit() // to flush the batch with the pruned heights
+		for v := int64(1); v < numVersions-int64(keepRecent); v++ {
+			if err := ms.LoadVersion(v); err == nil {
+				require.NoError(t, ms.LoadVersion(ls.Version)) // load latest
+				return false
+			}
+		}
+		return true
+	}
+	require.Eventually(t, isPruned, 1000*time.Second, 10*time.Millisecond, "expected error when loading pruned heights")
 }
 
 func TestMultiStore_PruningRestart(t *testing.T) {
@@ -618,10 +623,18 @@ func TestMultiStore_PruningRestart(t *testing.T) {
 	actualHeightToPrune = ms.pruningManager.GetPruningHeight(ms.LatestVersion())
 	require.Equal(t, int64(8), actualHeightToPrune)
 
-	for v := int64(1); v <= actualHeightToPrune; v++ {
-		_, err := ms.CacheMultiStoreWithVersion(v)
-		require.Error(t, err, "expected error when loading height: %d", v)
+	// Ensure async pruning is done
+	isPruned := func() bool {
+		ms.Commit() // to flush the batch with the pruned heights
+		for v := int64(1); v <= actualHeightToPrune; v++ {
+			if _, err := ms.CacheMultiStoreWithVersion(v); err == nil {
+				return false
+			}
+		}
+		return true
 	}
+
+	require.Eventually(t, isPruned, 1*time.Second, 10*time.Millisecond, "expected error when loading pruned heights")
 }
 
 // TestUnevenStoresHeightCheck tests if loading root store correctly errors when