From f40f820ec785f031326ace3f38c8ef13cde4e0e6 Mon Sep 17 00:00:00 2001
From: Julien Robert <julien@rbrt.fr>
Date: Tue, 21 Jan 2025 14:29:44 +0100
Subject: [PATCH] fix(x/protocolpool): withdraw rewards before export genesis
 (#23467)

---
 simapp/v2/app_config.go               | 28 +++++++++++++++++++++------
 x/protocolpool/keeper/genesis.go      | 15 ++++++++++++++
 x/protocolpool/keeper/genesis_test.go |  9 +++++++--
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/simapp/v2/app_config.go b/simapp/v2/app_config.go
index 9d9df139cbcb..7eb2f5ca68f5 100644
--- a/simapp/v2/app_config.go
+++ b/simapp/v2/app_config.go
@@ -154,7 +154,6 @@ var (
 						accounts.ModuleName,
 						authtypes.ModuleName,
 						banktypes.ModuleName,
-						// bankv2types.ModuleName,
 						distrtypes.ModuleName,
 						stakingtypes.ModuleName,
 						slashingtypes.ModuleName,
@@ -174,7 +173,28 @@ var (
 					},
 					// When ExportGenesis is not specified, the export genesis module order
 					// is equal to the init genesis order
-					// ExportGenesis: []string{},
+					ExportGenesis: []string{
+						consensustypes.ModuleName,
+						accounts.ModuleName,
+						authtypes.ModuleName,
+						pooltypes.ModuleName, // Must be exported before bank
+						banktypes.ModuleName,
+						distrtypes.ModuleName,
+						stakingtypes.ModuleName,
+						slashingtypes.ModuleName,
+						govtypes.ModuleName,
+						minttypes.ModuleName,
+						genutiltypes.ModuleName,
+						evidencetypes.ModuleName,
+						authz.ModuleName,
+						feegrant.ModuleName,
+						nft.ModuleName,
+						group.ModuleName,
+						upgradetypes.ModuleName,
+						vestingtypes.ModuleName,
+						circuittypes.ModuleName,
+						epochstypes.ModuleName,
+					},
 					// Uncomment if you want to set a custom migration order here.
 					// OrderMigrations: []string{},
 					// TODO GasConfig was added to the config in runtimev2.  Where/how was it set in v1?
@@ -298,10 +318,6 @@ var (
 				Name:   epochstypes.ModuleName,
 				Config: appconfig.WrapAny(&epochsmodulev1.Module{}),
 			},
-			// {
-			// 	Name:   bankv2types.ModuleName,
-			// 	Config: appconfig.WrapAny(&bankmodulev2.Module{}),
-			// },
 		},
 	}
 )
diff --git a/x/protocolpool/keeper/genesis.go b/x/protocolpool/keeper/genesis.go
index 65ae9537da48..1fb098823d7c 100644
--- a/x/protocolpool/keeper/genesis.go
+++ b/x/protocolpool/keeper/genesis.go
@@ -77,6 +77,21 @@ func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) error
 }
 
 func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error) {
+	// refresh all funds
+	if err := k.IterateAndUpdateFundsDistribution(ctx); err != nil {
+		return nil, err
+	}
+
+	// withdraw all rewards before exporting genesis
+	if err := k.RecipientFundDistribution.Walk(ctx, nil, func(key sdk.AccAddress, value types.DistributionAmount) (stop bool, err error) {
+		if _, err := k.withdrawRecipientFunds(ctx, key.Bytes()); err != nil {
+			return true, err
+		}
+		return false, nil
+	}); err != nil {
+		return nil, err
+	}
+
 	var cf []*types.ContinuousFund
 	err := k.ContinuousFund.Walk(ctx, nil, func(key sdk.AccAddress, value types.ContinuousFund) (stop bool, err error) {
 		recipient, err := k.authKeeper.AddressCodec().BytesToString(key)
diff --git a/x/protocolpool/keeper/genesis_test.go b/x/protocolpool/keeper/genesis_test.go
index 9e6c942a8da7..1a65369ad7c3 100644
--- a/x/protocolpool/keeper/genesis_test.go
+++ b/x/protocolpool/keeper/genesis_test.go
@@ -3,13 +3,18 @@ package keeper_test
 import (
 	"time"
 
+	"go.uber.org/mock/gomock"
+
 	"cosmossdk.io/math"
 	"cosmossdk.io/x/protocolpool/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
-func (suite *KeeperTestSuite) TestInitGenesis() {
+func (suite *KeeperTestSuite) TestInitExportGenesis() {
+	suite.bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), types.ProtocolPoolDistrAccount, gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	suite.bankKeeper.EXPECT().SendCoinsFromModuleToAccount(gomock.Any(), types.StreamAccount, gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+
 	hour := time.Hour
 	gs := types.NewGenesisState(
 		[]*types.ContinuousFund{
@@ -49,5 +54,5 @@ func (suite *KeeperTestSuite) TestInitGenesis() {
 	suite.Require().NoError(err)
 	suite.Require().Equal(gs.ContinuousFund, exportedGenState.ContinuousFund)
 	suite.Require().Equal(gs.Budget, exportedGenState.Budget)
-	suite.Require().Equal(math.NewInt(101), exportedGenState.LastBalance.Amount.AmountOf("stake"))
+	suite.Require().Equal(math.ZeroInt(), exportedGenState.LastBalance.Amount.AmountOf("stake"))
 }