From 49f3026d704c350526daecab01806ecbcacea779 Mon Sep 17 00:00:00 2001
From: Simon Noetzlin <simon.ntz@gmail.com>
Date: Wed, 9 Oct 2024 16:49:54 +0200
Subject: [PATCH] Introduce two following changes to the `commitBlock` method
 in testing: * increment the proposer priority of validators * update the
 proposer address in the current header

---
 testing/chain.go      |  5 ++++-
 testing/chain_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/testing/chain.go b/testing/chain.go
index bf45bd4d69e2..23752c987bf8 100644
--- a/testing/chain.go
+++ b/testing/chain.go
@@ -317,6 +317,9 @@ func (chain *TestChain) commitBlock(res *abci.ResponseFinalizeBlock) {
 	chain.Vals = chain.NextVals
 	chain.NextVals = ApplyValSetChanges(chain, chain.Vals, res.ValidatorUpdates)
 
+	// increment the proposer priority of validators
+	chain.Vals.IncrementProposerPriority(1)
+
 	// increment the current header
 	chain.CurrentHeader = cmtproto.Header{
 		ChainID: chain.ChainID,
@@ -327,7 +330,7 @@ func (chain *TestChain) commitBlock(res *abci.ResponseFinalizeBlock) {
 		Time:               chain.CurrentHeader.Time,
 		ValidatorsHash:     chain.Vals.Hash(),
 		NextValidatorsHash: chain.NextVals.Hash(),
-		ProposerAddress:    chain.CurrentHeader.ProposerAddress,
+		ProposerAddress:    chain.Vals.Proposer.Address,
 	}
 }
 
diff --git a/testing/chain_test.go b/testing/chain_test.go
index ea0e2cbe608a..21fa2d92f4e8 100644
--- a/testing/chain_test.go
+++ b/testing/chain_test.go
@@ -6,6 +6,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	sdkmath "cosmossdk.io/math"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 
@@ -41,3 +42,36 @@ func TestChangeValSet(t *testing.T) {
 	err = path.EndpointB.UpdateClient()
 	require.NoError(t, err)
 }
+
+func TestJailProposerValidator(t *testing.T) {
+	coord := ibctesting.NewCoordinator(t, 2)
+	chainA := coord.GetChain(ibctesting.GetChainID(1))
+	chainB := coord.GetChain(ibctesting.GetChainID(2))
+
+	path := ibctesting.NewPath(chainA, chainB)
+	coord.Setup(path)
+
+	// save valset length before jailing
+	valsetLen := len(chainA.Vals.Validators)
+
+	// jail the proposer validator in chain A
+	propAddr := sdk.ConsAddress(chainA.Vals.Proposer.Address)
+
+	err := chainA.GetSimApp().StakingKeeper.Jail(
+		chainA.GetContext(), propAddr)
+	require.NoError(t, err)
+
+	coord.CommitBlock(chainA)
+
+	// verify that update clients works even after validator update goes into effect
+	err = path.EndpointB.UpdateClient()
+	require.NoError(t, err)
+	err = path.EndpointB.UpdateClient()
+	require.NoError(t, err)
+
+	// check that the jailing has taken effect in chain A
+	require.Equal(t, valsetLen-1, len(chainA.Vals.Validators))
+
+	// check that the valset in chain A has a new proposer
+	require.False(t, propAddr.Equals(sdk.ConsAddress(chainA.Vals.Proposer.Address)))
+}