diff --git a/CHANGELOG.md b/CHANGELOG.md index f656130c3784..abf4bd83a231 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/simapp/state.go b/simapp/state.go index 836891cda88d..f491e098dd34 100644 --- a/simapp/state.go +++ b/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 diff --git a/types/simulation/types.go b/types/simulation/types.go index bdc0b59659d0..a152413f3c07 100644 --- a/types/simulation/types.go +++ b/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