feat(testutil/sims): add back AppStateFnWithExtendedCbs (#23236)

This commit is contained in:
mmsqe 2025-01-14 17:23:05 +08:00 committed by GitHub
parent b4e88cc517
commit 4466ea5121
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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 {