Skip to content

Commit

Permalink
feat: add moduleStateCb to allow access moduleState in sim test (#15903)
Browse files Browse the repository at this point in the history
Co-authored-by: Julien Robert <julien@rbrt.fr>
  • Loading branch information
mmsqe and julienrbrt authored Apr 24, 2023
1 parent 3030cc4 commit 4d41a87
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 15 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
* (x/consensus) [#15553](https://github.com/cosmos/cosmos-sdk/pull/15553) Migrate consensus module to use collections
* (x/bank) [#15764](https://github.com/cosmos/cosmos-sdk/pull/15764) Speedup x/bank InitGenesis
* (x/auth) [#15867](https://github.com/cosmos/cosmos-sdk/pull/15867) Support better logging for signature verification failure.
* (simtestutil) [#15903](https://github.com/cosmos/cosmos-sdk/pull/15903) Add `AppStateFnWithExtendedCbs` with moduleStateCb callback function to allow access moduleState.

### State Machine Breaking

Expand Down
40 changes: 30 additions & 10 deletions testutil/sims/state_helpers.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import (
"os"
"time"

"github.com/cosmos/gogoproto/proto"

"cosmossdk.io/math"

"github.com/cosmos/cosmos-sdk/codec"
Expand All @@ -29,23 +31,34 @@ const (
)

// AppStateFn 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.
// It calls AppStateFnWithExtendedCb with nil rawStateCb.
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 calls AppStateFnWithExtendedCbs with nil moduleStateCb.
func AppStateFnWithExtendedCb(
cdc codec.JSONCodec,
simManager *module.SimulationManager,
genesisState map[string]json.RawMessage,
rawStateCb func(rawState map[string]json.RawMessage),
) simtypes.AppStateFn {
return AppStateFnWithExtendedCbs(cdc, simManager, genesisState, nil, rawStateCb)
}

// 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.
// cb is the callback function to extend rawState.
func AppStateFnWithExtendedCb(
// moduleStateCb is the callback function to access moduleState.
// rawStateCb is the callback function to extend rawState.
func AppStateFnWithExtendedCbs(
cdc codec.JSONCodec,
simManager *module.SimulationManager,
genesisState map[string]json.RawMessage,
cb func(rawState map[string]json.RawMessage),
moduleStateCb func(moduleName string, genesisState interface{}),
rawStateCb func(rawState map[string]json.RawMessage),
) simtypes.AppStateFn {
return func(
r *rand.Rand,
Expand Down Expand Up @@ -148,12 +161,19 @@ func AppStateFnWithExtendedCb(
}

// change appState back
rawState[stakingtypes.ModuleName] = cdc.MustMarshalJSON(stakingState)
rawState[banktypes.ModuleName] = cdc.MustMarshalJSON(bankState)
for name, state := range map[string]proto.Message{
stakingtypes.ModuleName: stakingState,
banktypes.ModuleName: bankState,
} {
if moduleStateCb != nil {
moduleStateCb(name, state)
}
rawState[name] = cdc.MustMarshalJSON(state)
}

// extend state from callback function
if cb != nil {
cb(rawState)
if rawStateCb != nil {
rawStateCb(rawState)
}

// replace appstate
Expand Down
5 changes: 0 additions & 5 deletions types/simulation/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -177,11 +177,6 @@ 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

Expand Down

0 comments on commit 4d41a87

Please sign in to comment.