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

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 {