diff --git a/testutil/sims/state_helpers.go b/testutil/sims/state_helpers.go index 8e229891c7e7..53385e2c81d8 100644 --- a/testutil/sims/state_helpers.go +++ b/testutil/sims/state_helpers.go @@ -42,6 +42,23 @@ func AppStateFn( addressCodec, validatorCodec address.Codec, modules []module.AppModuleSimulation, genesisState map[string]json.RawMessage, +) simtypes.AppStateFn { + return AppStateFnWithExtendedCbs(cdc, addressCodec, validatorCodec, modules, genesisState, nil, nil) +} + +// 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. +// moduleStateCb is the callback function to access moduleState. +// postRawStateCb is the callback function to extend rawState. +func AppStateFnWithExtendedCbs( + cdc codec.JSONCodec, + addressCodec, validatorCodec address.Codec, + modules []module.AppModuleSimulation, + genesisState map[string]json.RawMessage, + moduleStateCb func(moduleName string, genesisState interface{}), + postRawStateCb func(rawState map[string]json.RawMessage), ) simtypes.AppStateFn { return func( r *rand.Rand, @@ -143,9 +160,17 @@ func AppStateFn( stakingtypes.ModuleName: stakingState, testutil.BankModuleName: bankState, } { + if moduleStateCb != nil { + moduleStateCb(name, state) + } rawState[name] = cdc.MustMarshalJSON(state) } + // extend state from callback function + if postRawStateCb != nil { + postRawStateCb(rawState) + } + // replace appstate appState, err = json.Marshal(rawState) if err != nil {