feat: add moduleStateCb to allow access moduleState in sim test (backport #15903)

CHANGELOG.md

@@ -40,6 +40,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### Improvements
 
 * (store) [#15683](https://github.com/cosmos/cosmos-sdk/pull/15683) `rootmulti.Store.CacheMultiStoreWithVersion` now can handle loading archival states that don't persist any of the module stores the current state has.
+* (simapp) [#15903](https://github.com/cosmos/cosmos-sdk/pull/15903) Add `AppStateFnWithExtendedCbs` with moduleStateCb callback function to allow access moduleState. Note, this function is present in simtestutil from `v0.47.2+`.
 
 ### Bug Fixes
 

simapp/state.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/gogo/protobuf/proto"
 	tmjson "github.com/tendermint/tendermint/libs/json"
 	tmtypes "github.com/tendermint/tendermint/types"
 
@@ -32,15 +33,28 @@ func AppStateFn(cdc codec.JSONCodec, simManager *module.SimulationManager) simty
 }
 
 // AppStateFnWithExtendedCb returns the initial application state using a genesis or the simulation parameters.
+// It calls AppStateFnWithExtendedCbs with nil moduleStateCb.
+func AppStateFnWithExtendedCb(
+	cdc codec.JSONCodec,
+	simManager *module.SimulationManager,
+	genesisState map[string]json.RawMessage,
+	rawStateCb func(rawState map[string]json.RawMessage),
+) simtypes.AppStateFn {
+	return AppStateFnWithExtendedCbs(cdc, simManager, genesisState, nil, rawStateCb)
+}
+
+// AppStateFnWithExtendedCbs returns the initial application state using a genesis or the simulation parameters.
 // It panics if the user provides files for both of them.
 // If a file is not given for the genesis or the sim params, it creates a randomized one.
-// genesisState is the genesis state of the app.
-// cb is the callback function to extend rawState.
-func AppStateFnWithExtendedCb(
+// genesisState is the default genesis state of the whole app.
+// moduleStateCb is the callback function to access moduleState.
+// rawStateCb is the callback function to extend rawState.
+func AppStateFnWithExtendedCbs(
 	cdc codec.JSONCodec,
 	simManager *module.SimulationManager,
 	genesisState map[string]json.RawMessage,
-	cb func(rawState map[string]json.RawMessage),
+	moduleStateCb func(moduleName string, genesisState interface{}),
+	rawStateCb func(rawState map[string]json.RawMessage),
 ) simtypes.AppStateFn {
 	return func(r *rand.Rand, accs []simtypes.Account, config simtypes.Config,
 	) (appState json.RawMessage, simAccs []simtypes.Account, chainID string, genesisTimestamp time.Time) {
@@ -139,12 +153,19 @@ func AppStateFnWithExtendedCb(
 		}
 
 		// change appState back
-		rawState[stakingtypes.ModuleName] = cdc.MustMarshalJSON(stakingState)
-		rawState[banktypes.ModuleName] = cdc.MustMarshalJSON(bankState)
+		for name, state := range map[string]proto.Message{
+			stakingtypes.ModuleName: stakingState,
+			banktypes.ModuleName:    bankState,
+		} {
+			if moduleStateCb != nil {
+				moduleStateCb(name, state)
+			}
+			rawState[name] = cdc.MustMarshalJSON(state)
+		}
 
 		// extend state from callback function
-		if cb != nil {
-			cb(rawState)
+		if rawStateCb != nil {
+			rawStateCb(rawState)
 		}
 
 		// replace appstate

types/simulation/types.go

@@ -162,9 +162,8 @@ type AppStateFn func(r *rand.Rand, accs []Account, config Config) (
 )
 
 // AppStateFnWithExtendedCb returns the app state json bytes and the genesis accounts
-type AppStateFnWithExtendedCb func(r *rand.Rand, accs []Account, config Config) (
-	appState json.RawMessage, accounts []Account, chainId string, genesisTimestamp time.Time,
-)
+// Deprecated: Use AppStateFn instead. This will be removed in a future relase.
+type AppStateFnWithExtendedCb AppStateFn
 
 // RandomAccountFn returns a slice of n random simulation accounts
 type RandomAccountFn func(r *rand.Rand, n int) []Account