From a794072e310d692f1a0ee3d68010de207bd8d5d7 Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Tue, 7 Jan 2025 21:47:31 +0800
Subject: [PATCH 1/2] feat(simsx): add back AppStateFnWithExtendedCbs

with moduleStateCb to access moduleState and rawStateCb to extend rawState
---
 testutil/sims/state_helpers.go | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/testutil/sims/state_helpers.go b/testutil/sims/state_helpers.go
index 76d06b8959d6..96d67d3a7e60 100644
--- a/testutil/sims/state_helpers.go
+++ b/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.
+// rawStateCb 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{}),
+	rawStateCb 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 rawStateCb != nil {
+			rawStateCb(rawState)
+		}
+
 		// replace appstate
 		appState, err = json.Marshal(rawState)
 		if err != nil {

From 17e9cb4b37e7e8ab28d82a7f31d2ebedf57a8ecf Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Wed, 8 Jan 2025 09:54:16 +0800
Subject: [PATCH 2/2] Apply suggestions from code review

---
 testutil/sims/state_helpers.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/testutil/sims/state_helpers.go b/testutil/sims/state_helpers.go
index 96d67d3a7e60..b433349b566f 100644
--- a/testutil/sims/state_helpers.go
+++ b/testutil/sims/state_helpers.go
@@ -51,14 +51,14 @@ func AppStateFn(
 // 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.
-// rawStateCb is the callback function to extend rawState.
+// 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{}),
-	rawStateCb func(rawState map[string]json.RawMessage),
+	postRawStateCb func(rawState map[string]json.RawMessage),
 ) simtypes.AppStateFn {
 	return func(
 		r *rand.Rand,
@@ -167,8 +167,8 @@ func AppStateFnWithExtendedCbs(
 		}
 
 		// extend state from callback function
-		if rawStateCb != nil {
-			rawStateCb(rawState)
+		if postRawStateCb != nil {
+			postRawStateCb(rawState)
 		}
 
 		// replace appstate