diff --git a/CHANGELOG.md b/CHANGELOG.md index 90ecc29180d3..c0d3d0831022 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,7 @@ Ref: https://keepachangelog.com/en/1.0.0/ # Changelog ## [Unreleased] +* (simapp) [#15305](https://github.com/cosmos/cosmos-sdk/pull/15305) Add `AppStateFnWithExtendedCb` with callback function to extend rawState. ## [v0.46.11](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.46.11) - 2022-03-03 diff --git a/simapp/state.go b/simapp/state.go index 135863c5bd7c..21e71470b24e 100644 --- a/simapp/state.go +++ b/simapp/state.go @@ -27,6 +27,14 @@ import ( // 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. func AppStateFn(cdc codec.JSONCodec, simManager *module.SimulationManager) simtypes.AppStateFn { + return AppStateFnWithExtendedCb(cdc, simManager, nil) +} + +// AppStateFnWithExtendedCb 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. +// cb is the callback function to extend rawState. +func AppStateFnWithExtendedCb(cdc codec.JSONCodec, simManager *module.SimulationManager, cb 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) { if FlagGenesisTimeValue == 0 { @@ -127,6 +135,11 @@ func AppStateFn(cdc codec.JSONCodec, simManager *module.SimulationManager) simty rawState[stakingtypes.ModuleName] = cdc.MustMarshalJSON(stakingState) rawState[banktypes.ModuleName] = cdc.MustMarshalJSON(bankState) + // extend state from callback function + if cb != nil { + cb(rawState) + } + // replace appstate appState, err = json.Marshal(rawState) if err != nil { diff --git a/types/simulation/types.go b/types/simulation/types.go index 199f47fa38e2..bdc0b59659d0 100644 --- a/types/simulation/types.go +++ b/types/simulation/types.go @@ -161,6 +161,11 @@ 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