From 544afb65e2d008fe2509f5be955d72170078f2f1 Mon Sep 17 00:00:00 2001
From: Dev Ojha <ValarDragon@users.noreply.github.com>
Date: Mon, 6 Jun 2022 08:17:48 -0500
Subject: [PATCH] refactor: Simplify SimulationManager setup (#12153)

---
 CHANGELOG.md               |  1 +
 simapp/app.go              | 19 ++++---------------
 types/module/simulation.go | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f8226fb79b..cd402996af 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -44,6 +44,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### Improvements
 
 * [#12089](https://github.com/cosmos/cosmos-sdk/pull/12089) Mark the `TipDecorator` as beta, don't include it in simapp by default.
+* [#12153](https://github.com/cosmos/cosmos-sdk/pull/12153) Add a new `NewSimulationManagerFromAppModules` constructor, to simplify simulation wiring.
 
 ### API Breaking Changes
 
diff --git a/simapp/app.go b/simapp/app.go
index 02d711bf1e..bae5e40657 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -373,21 +373,10 @@ func NewSimApp(
 	//
 	// NOTE: this is not required apps that don't use the simulator for fuzz testing
 	// transactions
-	app.sm = module.NewSimulationManager(
-		auth.NewAppModule(app.appCodec, app.AccountKeeper, authsims.RandomGenesisAccounts),
-		bank.NewAppModule(app.appCodec, app.BankKeeper, app.AccountKeeper),
-		feegrantmodule.NewAppModule(app.appCodec, app.AccountKeeper, app.BankKeeper, app.FeeGrantKeeper, app.interfaceRegistry),
-		gov.NewAppModule(app.appCodec, app.GovKeeper, app.AccountKeeper, app.BankKeeper),
-		mint.NewAppModule(app.appCodec, app.MintKeeper, app.AccountKeeper, nil),
-		staking.NewAppModule(app.appCodec, app.StakingKeeper, app.AccountKeeper, app.BankKeeper),
-		distr.NewAppModule(app.appCodec, app.DistrKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
-		slashing.NewAppModule(app.appCodec, app.SlashingKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
-		params.NewAppModule(app.ParamsKeeper),
-		evidence.NewAppModule(app.EvidenceKeeper),
-		authzmodule.NewAppModule(app.appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
-		groupmodule.NewAppModule(app.appCodec, app.GroupKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
-		nftmodule.NewAppModule(app.appCodec, app.NFTKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
-	)
+	overrideModules := map[string]module.AppModuleSimulation{
+		authtypes.ModuleName: auth.NewAppModule(app.appCodec, app.AccountKeeper, authsims.RandomGenesisAccounts),
+	}
+	app.sm = module.NewSimulationManagerFromAppModules(app.ModuleManager.Modules, overrideModules)
 
 	app.sm.RegisterStoreDecoders()
 
diff --git a/types/module/simulation.go b/types/module/simulation.go
index 591e11e084..8cba909345 100644
--- a/types/module/simulation.go
+++ b/types/module/simulation.go
@@ -3,6 +3,7 @@ package module
 import (
 	"encoding/json"
 	"math/rand"
+	"sort"
 	"time"
 
 	sdkmath "cosmossdk.io/math"
@@ -47,6 +48,38 @@ func NewSimulationManager(modules ...AppModuleSimulation) *SimulationManager {
 	}
 }
 
+// NewSimulationManagerFromAppModules creates a new SimulationManager object.
+//
+// First it sets any SimulationModule provided by overrideModules, and ignores any AppModule
+// with the same moduleName.
+// Then it attempts to cast every provided AppModule into an AppModuleSimulation.
+// If the cast succeeds, its included, otherwise it is excluded.
+func NewSimulationManagerFromAppModules(modules map[string]AppModule, overrideModules map[string]AppModuleSimulation) *SimulationManager {
+	simModules := []AppModuleSimulation{}
+	appModuleNamesSorted := make([]string, 0, len(modules))
+	for moduleName := range modules {
+		appModuleNamesSorted = append(appModuleNamesSorted, moduleName)
+	}
+
+	sort.Strings(appModuleNamesSorted)
+
+	for _, moduleName := range appModuleNamesSorted {
+		// for every module, see if we override it. If so, use override.
+		// Else, if we can cast the app module into a simulation module add it.
+		// otherwise no simulation module.
+		if simModule, ok := overrideModules[moduleName]; ok {
+			simModules = append(simModules, simModule)
+		} else {
+			appModule := modules[moduleName]
+			if simModule, ok := appModule.(AppModuleSimulation); ok {
+				simModules = append(simModules, simModule)
+			}
+			// cannot cast, so we continue
+		}
+	}
+	return NewSimulationManager(simModules...)
+}
+
 // GetProposalContents returns each module's proposal content generator function
 // with their default operation weight and key.
 func (sm *SimulationManager) GetProposalContents(simState SimulationState) []simulation.WeightedProposalContent {