From 7375a0676838f13fc3dc3ce116289ed20e0dc9f2 Mon Sep 17 00:00:00 2001
From: Robert Zaremba <robert@zaremba.ch>
Date: Fri, 18 Mar 2022 13:05:20 +0100
Subject: [PATCH] cleanup sim tests

---
 x/authz/keeper/keeper.go              |  1 +
 x/authz/simulation/operations.go      | 31 +++++++++++++--------------
 x/authz/simulation/operations_test.go | 12 ++++++-----
 3 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/x/authz/keeper/keeper.go b/x/authz/keeper/keeper.go
index 5708500a28a7..ff679211ba1d 100644
--- a/x/authz/keeper/keeper.go
+++ b/x/authz/keeper/keeper.go
@@ -239,6 +239,7 @@ func (k Keeper) GetAuthorizations(ctx sdk.Context, grantee sdk.AccAddress, grant
 // IterateGrants iterates over all authorization grants
 // This function should be used with caution because it can involve significant IO operations.
 // It should not be used in query or msg services without charging additional gas.
+// The iteration stops when the handler function returns true or the iterator exhaust.
 func (k Keeper) IterateGrants(ctx sdk.Context,
 	handler func(granterAddr sdk.AccAddress, granteeAddr sdk.AccAddress, grant authz.Grant) bool) {
 	store := ctx.KVStore(k.storeKey)
diff --git a/x/authz/simulation/operations.go b/x/authz/simulation/operations.go
index e7a0da7bdda6..a8eeef46d2a5 100644
--- a/x/authz/simulation/operations.go
+++ b/x/authz/simulation/operations.go
@@ -1,7 +1,6 @@
 package simulation
 
 import (
-	"fmt"
 	"math/rand"
 	"time"
 
@@ -216,19 +215,28 @@ func SimulateMsgExec(ak authz.AccountKeeper, bk authz.BankKeeper, k keeper.Keepe
 	return func(
 		r *rand.Rand, app *baseapp.BaseApp, ctx sdk.Context, accs []simtypes.Account, chainID string,
 	) (simtypes.OperationMsg, []simtypes.FutureOperation, error) {
-		hasGrant := false
-		var targetGrant authz.Grant
 		var granterAddr sdk.AccAddress
 		var granteeAddr sdk.AccAddress
+		var sendAuth *banktype.SendAuthorization
+		var err error
 		k.IterateGrants(ctx, func(granter, grantee sdk.AccAddress, grant authz.Grant) bool {
-			targetGrant = grant
 			granterAddr = granter
 			granteeAddr = grantee
-			hasGrant = true
-			return true
+			var a authz.Authorization
+			a, err = grant.GetAuthorization()
+			if err != nil {
+				return true
+			}
+			var ok bool
+			sendAuth, ok = a.(*banktype.SendAuthorization)
+			return ok
 		})
 
-		if !hasGrant {
+		if err != nil {
+			return simtypes.NoOpMsg(authz.ModuleName, TypeMsgExec, err.Error()), nil, err
+		}
+
+		if sendAuth == nil {
 			return simtypes.NoOpMsg(authz.ModuleName, TypeMsgExec, "no grant found"), nil, nil
 		}
 
@@ -250,15 +258,6 @@ func SimulateMsgExec(ak authz.AccountKeeper, bk authz.BankKeeper, k keeper.Keepe
 		}
 
 		msg := []sdk.Msg{banktype.NewMsgSend(granterAddr, granteeAddr, coins)}
-		authorization, err := targetGrant.GetAuthorization()
-		if err != nil {
-			return simtypes.NoOpMsg(authz.ModuleName, TypeMsgExec, err.Error()), nil, err
-		}
-
-		sendAuth, ok := authorization.(*banktype.SendAuthorization)
-		if !ok {
-			return simtypes.NoOpMsg(authz.ModuleName, TypeMsgExec, fmt.Sprintf("not a send authorization, got: %T", authorization)), nil, nil
-		}
 
 		_, err = sendAuth.Accept(ctx, msg[0])
 		if err != nil {
diff --git a/x/authz/simulation/operations_test.go b/x/authz/simulation/operations_test.go
index c8fd021fb14b..a050a5e10476 100644
--- a/x/authz/simulation/operations_test.go
+++ b/x/authz/simulation/operations_test.go
@@ -56,15 +56,17 @@ func (suite *SimTestSuite) TestWeightedOperations() {
 
 	require := suite.Require()
 	for i, w := range weightedOps {
-		operationMsg, _, err := w.Op()(r, suite.app.BaseApp, suite.ctx, accs, "")
+		op, _, err := w.Op()(r, suite.app.BaseApp, suite.ctx, accs, "")
 		require.NoError(err)
 		// the following checks are very much dependent from the ordering of the output given
 		// by WeightedOperations. if the ordering in WeightedOperations changes some tests
 		// will fail
-		require.Equal(expected[i].weight, w.Weight(), "weight should be the same")
-		require.Equal(expected[i].opMsgRoute, operationMsg.Route,
-			"route should be the same. %v", operationMsg.Comment)
-		require.Equal(expected[i].opMsgRoute, operationMsg.Name, "operation Msg name should be the same")
+		require.Equal(expected[i].weight, w.Weight(),
+			"weight should be the same. %v", op.Comment)
+		require.Equal(expected[i].opMsgRoute, op.Route,
+			"route should be the same. %v", op.Comment)
+		require.Equal(expected[i].opMsgRoute, op.Name,
+			"operation Msg name should be the same %v", op.Comment)
 	}
 }