feat: add extend cb to avoid unmarshal appState twice for sim test (#15305)

testutil/sims/state_helpers.go

@@ -33,6 +33,20 @@ const (
 // 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.
 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 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(
+	cdc codec.JSONCodec,
+	simManager *module.SimulationManager,
+	genesisState map[string]json.RawMessage,
+	cb func(rawState map[string]json.RawMessage),
+) simtypes.AppStateFn {
 	return func(
 		r *rand.Rand,
 		accs []simtypes.Account,
@@ -137,6 +151,11 @@ func AppStateFn(cdc codec.JSONCodec, simManager *module.SimulationManager, genes
 		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 {