test: fix key assignment in testing (#2344)

fix key assignment in testing

tests/integration/setup.go

@@ -9,20 +9,15 @@ import (
 	channeltypes "github.com/cosmos/ibc-go/v8/modules/core/04-channel/types"
 	ibctmtypes "github.com/cosmos/ibc-go/v8/modules/light-clients/07-tendermint"
 	ibctesting "github.com/cosmos/ibc-go/v8/testing"
-	"github.com/cosmos/ibc-go/v8/testing/mock"
 	"github.com/stretchr/testify/suite"
 
 	store "cosmossdk.io/store/types"
 
-	sdk "github.com/cosmos/cosmos-sdk/types"
-
 	abci "github.com/cometbft/cometbft/abci/types"
-	tmencoding "github.com/cometbft/cometbft/crypto/encoding"
 
 	icstestingutils "github.com/cosmos/interchain-security/v6/testutil/ibc_testing"
 	testutil "github.com/cosmos/interchain-security/v6/testutil/integration"
 	consumertypes "github.com/cosmos/interchain-security/v6/x/ccv/consumer/types"
-	providertypes "github.com/cosmos/interchain-security/v6/x/ccv/provider/types"
 	ccv "github.com/cosmos/interchain-security/v6/x/ccv/types"
 )
 
@@ -145,14 +140,6 @@ func (suite *CCVTestSuite) SetupTest() {
 	params.BlocksPerEpoch = 10
 	providerKeeper.SetParams(suite.providerCtx(), params)
 
-	// re-assign all validator keys for the first consumer chain
-	// this has to be done before:
-	//  1. the consumer chain is added to the coordinator
-	//  2. MakeGenesis is called on the provider chain
-	//  3. ibc/testing sets the tendermint header for the consumer chain app
-	providerKeeper.SetConsumerPhase(suite.providerCtx(), icstestingutils.FirstConsumerID, providertypes.CONSUMER_PHASE_INITIALIZED)
-	preProposalKeyAssignment(suite, icstestingutils.FirstConsumerID)
-
 	// start consumer chains
 	suite.consumerBundles = make(map[string]*icstestingutils.ConsumerBundle)
 	for i := 0; i < icstestingutils.NumConsumers; i++ {
@@ -416,34 +403,6 @@ func (s *CCVTestSuite) validateEndpointsClientConfig(consumerBundle icstestingut
 	)
 }
 
-// preProposalKeyAssignment assigns keys to all provider validators for
-// the consumer with consumerId before the chain is registered, i.e.,
-// before a client to the consumer is created
-func preProposalKeyAssignment(s *CCVTestSuite, consumerId string) {
-	providerKeeper := s.providerApp.GetProviderKeeper()
-
-	for _, val := range s.providerChain.Vals.Validators {
-		// get SDK validator
-		valAddr, err := sdk.ValAddressFromHex(val.Address.String())
-		s.Require().NoError(err)
-		validator := s.getVal(s.providerCtx(), valAddr)
-
-		// generate new PrivValidator
-		privVal := mock.NewPV()
-		tmPubKey, err := privVal.GetPubKey()
-		s.Require().NoError(err)
-		consumerKey, err := tmencoding.PubKeyToProto(tmPubKey)
-		s.Require().NoError(err)
-
-		// add Signer to the provider chain as there is no consumer chain to add it;
-		// as a result, NewTestChainWithValSet in AddConsumer uses providerChain.Signers
-		s.providerChain.Signers[tmPubKey.Address().String()] = privVal
-
-		err = providerKeeper.AssignConsumerKey(s.providerCtx(), consumerId, validator, consumerKey)
-		s.Require().NoError(err)
-	}
-}
-
 // packetSniffer implements the StreamingService interface.
 // Implements ListenEndBlock to record packets from events.
 type packetSniffer struct {

testutil/ibc_testing/generic_setup.go

@@ -7,6 +7,7 @@ import (
 
 	clienttypes "github.com/cosmos/ibc-go/v8/modules/core/02-client/types"
 	ibctesting "github.com/cosmos/ibc-go/v8/testing"
+	"github.com/cosmos/ibc-go/v8/testing/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
@@ -19,6 +20,7 @@ import (
 	testutil "github.com/cosmos/interchain-security/v6/testutil/integration"
 	testkeeper "github.com/cosmos/interchain-security/v6/testutil/keeper"
 	consumerkeeper "github.com/cosmos/interchain-security/v6/x/ccv/consumer/keeper"
+	providerkeeper "github.com/cosmos/interchain-security/v6/x/ccv/provider/keeper"
 	providertypes "github.com/cosmos/interchain-security/v6/x/ccv/provider/types"
 )
 
@@ -155,9 +157,6 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 	powerShapingParameters.Top_N = consumerTopNParams[index] // isn't used in CreateConsumerClient
 
 	consumerId := providerKeeper.FetchAndIncrementConsumerId(providerChain.GetContext())
-	if chainID == firstConsumerChainID {
-		FirstConsumerID = consumerId
-	}
 	providerKeeper.SetConsumerChainId(providerChain.GetContext(), consumerId, chainID)
 	err := providerKeeper.SetConsumerMetadata(providerChain.GetContext(), consumerId, consumerMetadata)
 	s.Require().NoError(err)
@@ -166,6 +165,15 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 	err = providerKeeper.SetConsumerPowerShapingParameters(providerChain.GetContext(), consumerId, powerShapingParameters)
 	s.Require().NoError(err)
 	providerKeeper.SetConsumerPhase(providerChain.GetContext(), consumerId, providertypes.CONSUMER_PHASE_INITIALIZED)
+	if chainID == firstConsumerChainID {
+		FirstConsumerID = consumerId
+		// re-assign all validator keys for the first consumer chain
+		// this has to be done before:
+		//  1. the consumer chain is added to the coordinator
+		//  2. MakeGenesis is called on the provider chain
+		//  3. ibc/testing sets the tendermint header for the consumer chain app
+		preProposalKeyAssignment(s, *providerChain, providerKeeper, providerApp, FirstConsumerID)
+	}
 	err = providerKeeper.AppendConsumerToBeLaunched(providerChain.GetContext(), consumerId, coordinator.CurrentTime)
 	s.Require().NoError(err)
 
@@ -228,3 +236,36 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 		TopN:       powerShapingParameters.Top_N,
 	}
 }
+
+// preProposalKeyAssignment assigns keys to all provider validators for
+// the consumer with consumerId before the chain is registered, i.e.,
+// before a client to the consumer is created
+func preProposalKeyAssignment(
+	s *suite.Suite,
+	providerChain ibctesting.TestChain,
+	providerKeeper providerkeeper.Keeper,
+	providerApp testutil.ProviderApp,
+	consumerId string,
+) {
+	for _, val := range providerChain.Vals.Validators {
+		// get SDK validator
+		valAddr, err := sdk.ValAddressFromHex(val.Address.String())
+		s.Require().NoError(err)
+		validator, err := providerApp.GetTestStakingKeeper().GetValidator(providerChain.GetContext(), valAddr)
+		s.Require().NoError(err)
+
+		// generate new PrivValidator
+		privVal := mock.NewPV()
+		tmPubKey, err := privVal.GetPubKey()
+		s.Require().NoError(err)
+		consumerKey, err := tmencoding.PubKeyToProto(tmPubKey)
+		s.Require().NoError(err)
+
+		// add Signer to the provider chain as there is no consumer chain to add it;
+		// as a result, NewTestChainWithValSet in AddConsumer uses providerChain.Signers
+		providerChain.Signers[tmPubKey.Address().String()] = privVal
+
+		err = providerKeeper.AssignConsumerKey(providerChain.GetContext(), consumerId, validator, consumerKey)
+		s.Require().NoError(err)
+	}
+}