fix: gov param change from swing (#68)

* working pre-remove validator

* fix test

* lint

e2e/helpers/gov.go

@@ -19,7 +19,7 @@ func ValidatorVote(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain,
 
 	resp, _ := cosmos.PollForProposalStatusV8(ctx, chain, height, height+searchHeightDelta, proposalID, cosmos.ProposalStatusPassedV8)
 	t.Log("PollForProposalStatusV8 resp", resp)
-	require.EqualValues(t, resp.Proposal.Status, cosmos.ProposalStatusPassedV8, "proposal status did not change to passed in expected number of blocks")
+	require.EqualValues(t, cosmos.ProposalStatusPassedV8, resp.Proposal.Status, "proposal status did not change to passed in expected number of blocks")
 }
 
 func SubmitParamChangeProp(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain, user ibc.Wallet, updatedParams []cosmosproto.Message, sender string, waitForBlocks uint64) string {

e2e/helpers/types.go

@@ -1,8 +1,6 @@
 package helpers
 
 import (
-	"time"
-
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/strangelove-ventures/poa"
 )
@@ -21,14 +19,14 @@ type Vals struct {
 		Description     struct {
 			Moniker string `json:"moniker"`
 		} `json:"description"`
-		UnbondingTime time.Time `json:"unbonding_time"`
+		UnbondingTime string `json:"unbonding_time"`
 		Commission    struct {
 			CommissionRates struct {
 				Rate          string `json:"rate"`
 				MaxRate       string `json:"max_rate"`
 				MaxChangeRate string `json:"max_change_rate"`
 			} `json:"commission_rates"`
-			UpdateTime time.Time `json:"update_time"`
+			UpdateTime string `json:"update_time"`
 		} `json:"commission"`
 		MinSelfDelegation string `json:"min_self_delegation"`
 	} `json:"validators"`
@@ -44,9 +42,9 @@ type BlockData struct {
 			Block string `json:"block"`
 			App   string `json:"app"`
 		} `json:"version"`
-		ChainID     string    `json:"chain_id"`
-		Height      string    `json:"height"`
-		Time        time.Time `json:"time"`
+		ChainID     string `json:"chain_id"`
+		Height      string `json:"height"`
+		Time        string `json:"time"`
 		LastBlockID struct {
 			Hash          string `json:"hash"`
 			PartSetHeader struct {
@@ -81,10 +79,10 @@ type BlockData struct {
 			} `json:"part_set_header"`
 		} `json:"block_id"`
 		Signatures []struct {
-			BlockIDFlag      string    `json:"block_id_flag"`
-			ValidatorAddress string    `json:"validator_address"`
-			Timestamp        time.Time `json:"timestamp"`
-			Signature        string    `json:"signature"`
+			BlockIDFlag      string `json:"block_id_flag"`
+			ValidatorAddress string `json:"validator_address"`
+			Timestamp        string `json:"timestamp"`
+			Signature        string `json:"signature"`
 		} `json:"signatures"`
 	} `json:"last_commit"`
 }

e2e/poa_test.go

@@ -54,10 +54,9 @@ func TestPOA(t *testing.T) {
 	// === Test Cases ===
 	testStakingDisabled(t, ctx, chain, validators, acc0)
 	testGovernance(t, ctx, chain, acc0, validators)
-	testUpdatePOAParams(t, ctx, chain, validators, acc0, incorrectUser)
 	testPowerErrors(t, ctx, chain, validators, incorrectUser, acc0)
 	testRemoveValidator(t, ctx, chain, validators, acc0)
-	testPowerSwing(t, ctx, chain, validators, acc0)
+	testUpdatePOAParams(t, ctx, chain, validators, acc0, incorrectUser)
 }
 
 func testUpdatePOAParams(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain, validators []string, acc0, incorrectUser ibc.Wallet) {
@@ -180,15 +179,6 @@ func testRemoveValidator(t *testing.T, ctx context.Context, chain *cosmos.Cosmos
 	assertSignatures(t, ctx, chain, 1)
 }
 
-func testPowerSwing(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain, validators []string, acc0 ibc.Wallet) {
-	t.Log("\n===== TEST POWER SWING =====")
-	helpers.POASetPower(t, ctx, chain, acc0, validators[0], 1_000_000_000_000_000, "--unsafe")
-	if err := testutil.WaitForBlocks(ctx, 2, chain); err != nil {
-		t.Fatal(err)
-	}
-	helpers.POASetPower(t, ctx, chain, acc0, validators[0], 1_000_000, "--unsafe")
-}
-
 func testStakingDisabled(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain, validators []string, acc0 ibc.Wallet) {
 	t.Log("\n===== TEST STAKING DISABLED =====")
 	txRes, _ := helpers.StakeTokens(t, ctx, chain, acc0, validators[0], "1stake")

e2e/setup.go

@@ -8,7 +8,7 @@ import (
 )
 
 var (
-	VotingPeriod     = "30s"
+	VotingPeriod     = "15s"
 	MaxDepositPeriod = "10s"
 	Denom            = "stake"
 

keeper/keeper_test.go

@@ -155,15 +155,15 @@ func (f *testFixture) createBaseStakingValidators(t *testing.T) {
 		_, err := f.msgServer.SetPower(f.ctx, &poa.MsgSetPower{
 			Sender:           f.addrs[0].String(),
 			ValidatorAddress: valAddr,
-			Power:            1_000_000,
+			Power:            bondCoin.Amount.Uint64(),
 			Unsafe:           true,
 		})
 		require.NoError(t, err)
 
 		// increase the block so the new validator is in the validator set
-		f.ctx = f.ctx.WithBlockHeight(f.ctx.BlockHeight() + 1)
-		_, err = f.stakingKeeper.ApplyAndReturnValidatorSetUpdates(f.ctx)
+		_, err = f.IncreaseBlock(1)
 		require.NoError(t, err)
+		// fmt.Printf("createBaseStakingValidators updates : %+v", updates)
 
 		valAddrBz, err := sdk.ValAddressFromBech32(val.GetOperator())
 		require.NoError(t, err)
@@ -175,7 +175,19 @@ func (f *testFixture) createBaseStakingValidators(t *testing.T) {
 		if err := f.stakingKeeper.SetValidator(f.ctx, validator); err != nil {
 			panic(err)
 		}
+	}
+
+	// update validator set power
+	allVals, err := f.stakingKeeper.GetValidators(f.ctx, 100)
+	require.NoError(t, err)
 
+	totalPower := math.ZeroInt()
+	for _, val := range allVals {
+		totalPower = totalPower.Add(val.Tokens)
+	}
+
+	if err := f.stakingKeeper.SetLastTotalPower(f.ctx, totalPower); err != nil {
+		panic(err)
 	}
 }
 

keeper/msg_server.go

@@ -99,6 +99,10 @@ func (ms msgServer) RemoveValidator(ctx context.Context, msg *poa.MsgRemoveValid
 		return nil, err
 	}
 
+	if _, err := ms.updatePOAPower(ctx, msg.ValidatorAddress, 0); err != nil {
+		return nil, err
+	}
+
 	return &poa.MsgRemoveValidatorResponse{}, nil
 }
 
@@ -331,9 +335,36 @@ func (ms msgServer) updatePOAPower(ctx context.Context, valOpBech32 string, powe
 		return stakingtypes.Validator{}, err
 	}
 
+	// Update the total power cache
+	totalSetPower, err := ms.getValidatorSetTotalPower(ctx)
+	if err != nil {
+		return stakingtypes.Validator{}, err
+	}
+
+	if err := ms.k.stakingKeeper.SetLastTotalPower(ctx, totalSetPower); err != nil {
+		return stakingtypes.Validator{}, err
+	}
+
 	return val, nil
 }
 
+func (ms msgServer) getValidatorSetTotalPower(ctx context.Context) (math.Int, error) {
+	vals, err := ms.k.stakingKeeper.GetAllValidators(ctx)
+	if err != nil {
+		return math.ZeroInt(), err
+	}
+
+	totalPower := math.ZeroInt()
+	for _, val := range vals {
+		if val.Status != stakingtypes.Bonded {
+			continue
+		}
+		totalPower = totalPower.Add(val.Tokens)
+	}
+
+	return totalPower, nil
+}
+
 func (ms msgServer) isAdmin(ctx context.Context, fromAddr string) bool {
 	for _, auth := range ms.k.GetAdmins(ctx) {
 		if auth == fromAddr {

keeper/msg_server_test.go

@@ -278,6 +278,77 @@ func TestRemoveValidator(t *testing.T) {
 	}
 }
 
+func TestGlobalPowerUpdates(t *testing.T) {
+	f := SetupTest(t)
+	require := require.New(t)
+
+	vals, err := f.stakingKeeper.GetValidators(f.ctx, 100)
+	require.NoError(err)
+
+	totalValTokens := math.ZeroInt()
+	for _, val := range vals {
+		totalValTokens = totalValTokens.Add(val.Tokens)
+	}
+
+	testCases := []struct {
+		name               string
+		createNewValidator bool
+		request            *poa.MsgSetPower
+		expectErrMsg       string
+	}{
+		{
+			name:               "new validator swing",
+			createNewValidator: true,
+			request: &poa.MsgSetPower{
+				Sender: f.addrs[0].String(),
+				Power:  2_000_000,
+				Unsafe: true,
+			},
+			expectErrMsg: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+
+			// increase block
+			if _, err := f.IncreaseBlock(1); err != nil {
+				panic(err)
+			}
+
+			// get total set power
+			total, err := f.stakingKeeper.GetLastTotalPower(f.ctx)
+			require.NoError(err)
+			require.EqualValues(totalValTokens, total)
+
+			// add a new validator if the test case requires it
+			if tc.createNewValidator {
+				valAddr := f.CreatePendingValidator(fmt.Sprintf("val-%s", tc.name), tc.request.Power)
+				tc.request.ValidatorAddress = valAddr.String()
+
+				// check the pending validators includes the new validator
+				pendingVals, err := f.k.GetPendingValidators(f.ctx)
+				require.NoError(err)
+				require.EqualValues(1, len(pendingVals.Validators))
+			}
+
+			_, err = f.msgServer.SetPower(f.ctx, tc.request)
+			if tc.expectErrMsg != "" {
+				require.Error(err)
+				require.ErrorContains(err, tc.expectErrMsg)
+			} else {
+				require.NoError(err)
+			}
+
+			// verify that it goes up the Power request amount
+			total, err = f.stakingKeeper.GetLastTotalPower(f.ctx)
+			require.NoError(err)
+			require.EqualValues(totalValTokens.AddRaw(int64(tc.request.Power)), total)
+		})
+	}
+}
+
 // mintTokensToBondedPool mints tokens to the bonded pool so the validator set
 // in testing can be removed.
 // In the future, this same logic would be run during the migration from POA->POS.

module/abci.go

@@ -3,6 +3,7 @@ package module
 import (
 	"context"
 
+	"cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 )
@@ -13,6 +14,12 @@ func (am AppModule) BeginBlocker(ctx context.Context) error {
 		return err
 	}
 
+	if sdk.UnwrapSDKContext(ctx).BlockHeight() == 2 {
+		if err := am.updateLastPower(ctx, vals); err != nil {
+			return err
+		}
+	}
+
 	for _, v := range vals {
 		switch v.GetStatus() {
 
@@ -21,6 +28,7 @@ func (am AppModule) BeginBlocker(ctx context.Context) error {
 			if err := am.keeper.GetStakingKeeper().SetValidator(ctx, v); err != nil {
 				return err
 			}
+
 		case stakingtypes.Unbonded:
 			valAddr, err := sdk.ValAddressFromBech32(v.OperatorAddress)
 			if err != nil {
@@ -35,3 +43,25 @@ func (am AppModule) BeginBlocker(ctx context.Context) error {
 
 	return nil
 }
+
+func (am AppModule) updateLastPower(ctx context.Context, vals []stakingtypes.Validator) error {
+	totalPower := math.ZeroInt()
+	for _, v := range vals {
+		totalPower = totalPower.Add(v.Tokens)
+
+		valAddr, err := sdk.ValAddressFromBech32(v.OperatorAddress)
+		if err != nil {
+			return err
+		}
+
+		if err := am.keeper.GetStakingKeeper().SetLastValidatorPower(ctx, valAddr, v.Tokens.Int64()); err != nil {
+			return err
+		}
+	}
+
+	if err := am.keeper.GetStakingKeeper().SetLastTotalPower(ctx, totalPower); err != nil {
+		return err
+	}
+
+	return nil
+}