feat: add extend cb to avoid unmarshal state twice for sim test (back…

…port: cosmos#15305) (cosmos#15321)

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
 

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 {

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