From 053ae2687c8ea1b82942456567a3bf705f9372a6 Mon Sep 17 00:00:00 2001
From: sontrinh16 <trinhleson2000@gmail.com>
Date: Tue, 11 Jun 2024 16:59:01 +0700
Subject: [PATCH] address comment

---
 UPGRADING.md            |  9 ++-------
 simapp/upgrades.go      | 13 +------------
 simapp/upgrades_test.go |  2 +-
 x/auth/keeper/keeper.go | 15 +++++++++++++++
 4 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/UPGRADING.md b/UPGRADING.md
index fb3634195230..2709df650c97 100644
--- a/UPGRADING.md
+++ b/UPGRADING.md
@@ -254,15 +254,10 @@ For modules that have migrated, verify you are checking against `collections.Err
 
 #### `x/accounts`
 
-Accounts's AccountNumber will be used as a global account number tracking replacing Auth legacy AccountNumber. Must set accounts's AccountNumber with auth's AccountNumber value in upgrade handler:
+Accounts's AccountNumber will be used as a global account number tracking replacing Auth legacy AccountNumber. Must set accounts's AccountNumber with auth's AccountNumber value in upgrade handler. This is done through auth keeper MigrateAccountNumber function.
 
 ```go
-currentAccNum, err := app.AuthKeeper.RemoveLegacyAccountNumber(ctx)
-if err != nil {
-	return nil, err
-}
-
-err = app.AccountsKeeper.InitAccountNumberSeqUnsafe(ctx, currentAccNum)
+err := app.AuthKeeper.MigrateAccountNumber(ctx)
 if err != nil {
 	return nil, err
 }
diff --git a/simapp/upgrades.go b/simapp/upgrades.go
index 820848d3decf..38206fb8360c 100644
--- a/simapp/upgrades.go
+++ b/simapp/upgrades.go
@@ -26,7 +26,7 @@ func (app SimApp) RegisterUpgradeHandlers() {
 		UpgradeName,
 		func(ctx context.Context, _ upgradetypes.Plan, fromVM appmodule.VersionMap) (appmodule.VersionMap, error) {
 			// sync accounts and auth module account number
-			err := app.syncAccountNumber(ctx)
+			err := app.AuthKeeper.MigrateAccountNumberUnsafe(ctx)
 			if err != nil {
 				return nil, err
 			}
@@ -56,14 +56,3 @@ func (app SimApp) RegisterUpgradeHandlers() {
 		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
 	}
 }
-
-func (app SimApp) syncAccountNumber(ctx context.Context) error {
-	currentAccNum, err := app.AuthKeeper.RemoveLegacyAccountNumber(ctx)
-	if err != nil {
-		return err
-	}
-
-	err = app.AccountsKeeper.InitAccountNumberSeqUnsafe(ctx, currentAccNum)
-
-	return err
-}
diff --git a/simapp/upgrades_test.go b/simapp/upgrades_test.go
index c569ef25f24a..3b34120d791d 100644
--- a/simapp/upgrades_test.go
+++ b/simapp/upgrades_test.go
@@ -37,7 +37,7 @@ func TestSyncAccountNumber(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, uint64(10), num)
 
-	app.syncAccountNumber(ctx)
+	app.AuthKeeper.MigrateAccountNumberUnsafe(ctx)
 
 	// make sure the DB entry for this key is deleted
 	v, err = store.Get(bytesKey)
diff --git a/x/auth/keeper/keeper.go b/x/auth/keeper/keeper.go
index 6f2adc895bb3..c70f6c7d2c3d 100644
--- a/x/auth/keeper/keeper.go
+++ b/x/auth/keeper/keeper.go
@@ -338,3 +338,18 @@ func (ak AccountKeeper) NonAtomicMsgsExec(ctx context.Context, signer sdk.AccAdd
 
 	return msgResponses, nil
 }
+
+// MigrateAccountNumberUnsafe migrates auth's account number to accounts's account number
+// and delete store entry for auth legacy GlobalAccountNumberKey.
+//
+// Should only use in an upgrade handler for migrate process.
+func (ak AccountKeeper) MigrateAccountNumberUnsafe(ctx context.Context) error {
+	currentAccNum, err := ak.RemoveLegacyAccountNumber(ctx)
+	if err != nil {
+		return err
+	}
+
+	err = ak.AccountsModKeeper.InitAccountNumberSeqUnsafe(ctx, currentAccNum)
+
+	return err
+}