From 79a6cd8e3b77a59176d80b2838637f2ae3bb12c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E5=BF=97=E5=BC=BA?= <zhiqiang@bianjie.ai>
Date: Mon, 16 Nov 2020 10:26:56 +0800
Subject: [PATCH] cherry-pick from
 https://github.com/cosmos/cosmos-sdk/pull/7817

---
 store/rootmulti/store.go            |  4 ++-
 x/upgrade/types/storeloader.go      |  6 ++---
 x/upgrade/types/storeloader_test.go | 42 +++++++++++++++++------------
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index b0ca3010baef..e026eb34dffd 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -320,7 +320,9 @@ func (rs *Store) TracingEnabled() bool {
 // LastCommitID implements Committer/CommitStore.
 func (rs *Store) LastCommitID() types.CommitID {
 	if rs.lastCommitInfo == nil {
-		return types.CommitID{}
+		return types.CommitID{
+			Version: getLatestVersion(rs.db),
+		}
 	}
 
 	return rs.lastCommitInfo.CommitID()
diff --git a/x/upgrade/types/storeloader.go b/x/upgrade/types/storeloader.go
index 7cec2d4b610f..4bc513f14537 100644
--- a/x/upgrade/types/storeloader.go
+++ b/x/upgrade/types/storeloader.go
@@ -2,7 +2,6 @@ package types
 
 import (
 	"github.com/cosmos/cosmos-sdk/baseapp"
-	"github.com/cosmos/cosmos-sdk/store/rootmulti"
 	store "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -11,11 +10,10 @@ import (
 // pattern. This is useful for custom upgrade loading logic.
 func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades) baseapp.StoreLoader {
 	return func(ms sdk.CommitMultiStore) error {
-		latestVersion := ms.(*rootmulti.Store).GetLatestVersion()
-		if upgradeHeight == latestVersion+1 {
+		if upgradeHeight == ms.LastCommitID().Version+1 {
 			// Check if the current commit version and upgrade height matches
 			if len(storeUpgrades.Renamed) > 0 ||
-				len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 {
+				len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 || len(storeUpgrades.Added) > 0 {
 				return ms.LoadLatestVersionAndUpgrade(storeUpgrades)
 			}
 		}
diff --git a/x/upgrade/types/storeloader_test.go b/x/upgrade/types/storeloader_test.go
index 2814b61b4c1c..ec2bfa824d07 100644
--- a/x/upgrade/types/storeloader_test.go
+++ b/x/upgrade/types/storeloader_test.go
@@ -65,11 +65,13 @@ func checkStore(t *testing.T, db dbm.DB, ver int64, storeKey string, k, v []byte
 // Test that we can make commits and then reload old versions.
 // Test that LoadLatestVersion actually does.
 func TestSetLoader(t *testing.T) {
+	upgradeHeight := int64(5)
+
 	// set a temporary home dir
 	homeDir := t.TempDir()
 	upgradeInfoFilePath := filepath.Join(homeDir, "upgrade-info.json")
 	upgradeInfo := &store.UpgradeInfo{
-		Name: "test", Height: 0,
+		Name: "test", Height: upgradeHeight,
 	}
 
 	data, err := json.Marshal(upgradeInfo)
@@ -82,8 +84,6 @@ func TestSetLoader(t *testing.T) {
 	_, err = os.Stat(upgradeInfoFilePath)
 	require.NoError(t, err)
 
-	upgradeHeight := int64(2)
-
 	cases := map[string]struct {
 		setLoader    func(*baseapp.BaseApp)
 		origStoreKey string
@@ -113,33 +113,41 @@ func TestSetLoader(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			// prepare a db with some data
 			db := dbm.NewMemDB()
+
 			initStore(t, db, tc.origStoreKey, k, v)
 
 			// load the app with the existing db
 			opts := []func(*baseapp.BaseApp){baseapp.SetPruning(store.PruneNothing)}
+
+			origapp := baseapp.NewBaseApp(t.Name(), defaultLogger(), db, nil, opts...)
+			origapp.MountStores(sdk.NewKVStoreKey(tc.origStoreKey))
+			err := origapp.LoadLatestVersion()
+			require.Nil(t, err)
+
+			for i := int64(2); i <= upgradeHeight-1; i++ {
+				origapp.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: i}})
+				res := origapp.Commit()
+				require.NotNil(t, res.Data)
+			}
+
 			if tc.setLoader != nil {
 				opts = append(opts, tc.setLoader)
 			}
 
+			// load the new app with the original app db
 			app := baseapp.NewBaseApp(t.Name(), defaultLogger(), db, nil, opts...)
-			capKey := sdk.NewKVStoreKey("main")
-			app.MountStores(capKey)
 			app.MountStores(sdk.NewKVStoreKey(tc.loadStoreKey))
-			err := app.LoadLatestVersion()
+			err = app.LoadLatestVersion()
 			require.Nil(t, err)
 
-			for i := app.LastBlockHeight() + 1; i <= upgradeHeight * 2; i++ {
-				// "execute" one block
-				app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: i}})
-				res := app.Commit()
-				require.NotNil(t, res.Data)
+			// "execute" one block
+			app.BeginBlock(abci.RequestBeginBlock{Header: tmproto.Header{Height: upgradeHeight}})
+			res := app.Commit()
+			require.NotNil(t, res.Data)
 
-				if i >= upgradeHeight {
-					// check db is properly updated
-					checkStore(t, db, i, tc.loadStoreKey, k, v)
-					checkStore(t, db, i, tc.loadStoreKey, []byte("foo"), nil)
-				}
-			}
+			// check db is properly updated
+			checkStore(t, db, upgradeHeight, tc.loadStoreKey, k, v)
+			checkStore(t, db, upgradeHeight, tc.loadStoreKey, []byte("foo"), nil)
 		})
 	}
 }