diff --git a/CHANGELOG.md b/CHANGELOG.md index 80175c5dee45..c3360d8ae415 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -102,6 +102,7 @@ Ref: https://keepachangelog.com/en/1.0.0/ * (x/consensus) [#15553](https://github.com/cosmos/cosmos-sdk/pull/15553) Migrate consensus module to use collections * (x/bank) [#15764](https://github.com/cosmos/cosmos-sdk/pull/15764) Speedup x/bank InitGenesis * (x/auth) [#15867](https://github.com/cosmos/cosmos-sdk/pull/15867) Support better logging for signature verification failure. +* (simtestutil) [#15903](https://github.com/cosmos/cosmos-sdk/pull/15903) Add `AppStateFnWithExtendedCbs` with moduleStateCb callback function to allow access moduleState. ### State Machine Breaking diff --git a/testutil/sims/state_helpers.go b/testutil/sims/state_helpers.go index 7b97cb6160cf..73319ecaba48 100644 --- a/testutil/sims/state_helpers.go +++ b/testutil/sims/state_helpers.go @@ -8,6 +8,8 @@ import ( "os" "time" + "github.com/cosmos/gogoproto/proto" + "cosmossdk.io/math" "github.com/cosmos/cosmos-sdk/codec" @@ -29,23 +31,34 @@ const ( ) // AppStateFn 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 default genesis state of the whole app. +// It calls AppStateFnWithExtendedCb with nil rawStateCb. func AppStateFn(cdc codec.JSONCodec, simManager *module.SimulationManager, genesisState map[string]json.RawMessage) simtypes.AppStateFn { return AppStateFnWithExtendedCb(cdc, simManager, genesisState, nil) } // 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 default genesis state of the whole app. -// cb is the callback function to extend rawState. -func AppStateFnWithExtendedCb( +// 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, @@ -148,12 +161,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 56d496e03b0e..34d29d04e323 100644 --- a/types/simulation/types.go +++ b/types/simulation/types.go @@ -177,11 +177,6 @@ type AppStateFn func(r *rand.Rand, accs []Account, config Config) ( appState json.RawMessage, accounts []Account, chainId string, genesisTimestamp time.Time, ) -// 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, -) - // RandomAccountFn returns a slice of n random simulation accounts type RandomAccountFn func(r *rand.Rand, n int) []Account