Add SudoContractProposal simulation

CosmWasm · Nov 25, 2022 · 9a785e8 · 9a785e8
1 parent 1b8cf7b
commit 9a785e8
Showing 1 changed file with 35 additions and 19 deletions.
diff --git a/x/wasm/simulation/proposals.go b/x/wasm/simulation/proposals.go
@@ -66,15 +66,24 @@ func ProposalContents(bk BankKeeper, wasmKeeper WasmKeeper) []simtypes.WeightedP
 			params.DefaultWeightClearAdminProposal,
 			SimulateClearAdminProposal(
 				wasmKeeper,
-				DefaultSimulateClearAdminProposalContractSelector,
+				DefaultSimulateContractSelector,
 			),
 		),
 		simulation.NewWeightedProposalContent(
 			WeightMigrateContractProposal,
 			params.DefaultWeightMigrateContractProposal,
 			SimulateMigrateContractProposal(
 				wasmKeeper,
-				DefaultSimulateMigrateContractProposalContractSelector,
+				DefaultSimulateContractSelector,
+				DefaultSimulationCodeIDSelector,
+			),
+		),
+		simulation.NewWeightedProposalContent(
+			WeightSudoContractProposal,
+			params.DefaultWeightSudoContractProposal,
+			SimulateSudoContractProposal(
+				wasmKeeper,
+				DefaultSimulateContractSelector,
 			),
 		),
 	}
@@ -210,7 +219,7 @@ func SimulateUpdateAdminProposal(wasmKeeper WasmKeeper, contractSelector UpdateA
 
 type ClearAdminContractSelector func(sdk.Context, WasmKeeper) sdk.AccAddress
 
-func DefaultSimulateClearAdminProposalContractSelector(
+func DefaultSimulateContractSelector(
 	ctx sdk.Context,
 	wasmKeeper WasmKeeper,
 ) sdk.AccAddress {
@@ -240,36 +249,43 @@ func SimulateClearAdminProposal(wasmKeeper WasmKeeper, contractSelector ClearAdm
 
 type MigrateContractProposalContractSelector func(sdk.Context, WasmKeeper) sdk.AccAddress
 
-func DefaultSimulateMigrateContractProposalContractSelector(
-	ctx sdk.Context,
-	wasmKeeper WasmKeeper,
-) sdk.AccAddress {
-	var contractAddr sdk.AccAddress
-	wasmKeeper.IterateContractInfo(ctx, func(address sdk.AccAddress, info types.ContractInfo) bool {
-		contractAddr = address
-		return true
-	})
-	return contractAddr
-}
-
 // Simulate migrate contract proposal
-func SimulateMigrateContractProposal(wasmKeeper WasmKeeper, contractSelector MigrateContractProposalContractSelector) simtypes.ContentSimulatorFn {
+func SimulateMigrateContractProposal(wasmKeeper WasmKeeper, contractSelector MigrateContractProposalContractSelector, codeSelector CodeIDSelector) simtypes.ContentSimulatorFn {
 	return func(r *rand.Rand, ctx sdk.Context, accs []simtypes.Account) simtypes.Content {
 		ctAddress := contractSelector(ctx, wasmKeeper)
 		if ctAddress == nil {
 			return nil
 		}
 
-		codeID, err := simtypes.RandPositiveInt(r, sdk.NewInt(1000))
-		if err != nil {
+		codeID := codeSelector(ctx, wasmKeeper)
+		if codeID == 0 {
 			return nil
 		}
 
 		return types.NewMigrateContractProposal(
 			simtypes.RandStringOfLength(r, 10),
 			simtypes.RandStringOfLength(r, 10),
 			ctAddress.String(),
-			codeID.Uint64(),
+			codeID,
+			[]byte(`{}`),
+		)
+	}
+}
+
+type SudoContractProposalContractSelector func(sdk.Context, WasmKeeper) sdk.AccAddress
+
+// Simulate sudo contract proposal
+func SimulateSudoContractProposal(wasmKeeper WasmKeeper, contractSelector SudoContractProposalContractSelector) simtypes.ContentSimulatorFn {
+	return func(r *rand.Rand, ctx sdk.Context, accs []simtypes.Account) simtypes.Content {
+		ctAddress := contractSelector(ctx, wasmKeeper)
+		if ctAddress == nil {
+			return nil
+		}
+
+		return types.NewSudoContractProposal(
+			simtypes.RandStringOfLength(r, 10),
+			simtypes.RandStringOfLength(r, 10),
+			ctAddress.String(),
 			[]byte(`{}`),
 		)
 	}