test: update integration test suite for PSS

tests/integration/distribution.go

@@ -979,3 +979,95 @@
 		})
 	}
 }
+
+// TestMultiConsumerRewardsDistribution tests the rewards distribution of multiple consumers chains
+func (s *CCVTestSuite) TestMultiConsumerRewardsDistribution() {
+	s.SetupAllCCVChannels()
+	s.SetupAllTransferChannels()
+
+	providerBankKeeper := s.providerApp.GetTestBankKeeper()
+	providerAccountKeeper := s.providerApp.GetTestAccountKeeper()
+
+	// check that the reward provider pool is empty
+	rewardPool := providerAccountKeeper.GetModuleAccount(s.providerCtx(), providertypes.ConsumerRewardsPool).GetAddress()
+	rewardCoins := providerBankKeeper.GetAllBalances(s.providerCtx(), rewardPool)
+	s.Require().Empty(rewardCoins)
+
+	totalConsumerRewards := sdk.Coins{}
+
+	// Iterate over the consumers and perform the reward distribution
+	// to the provider
+	for chainID := range s.consumerBundles {
+		bundle := s.consumerBundles[chainID]
+		consumerKeeper := bundle.App.GetConsumerKeeper()
+		bankKeeper := bundle.App.GetTestBankKeeper()
+		accountKeeper := bundle.App.GetTestAccountKeeper()
+
+		// set the consumer reward denom and the block per distribution params
+		params := consumerKeeper.GetConsumerParams(bundle.GetCtx())
+		params.RewardDenoms = []string{sdk.DefaultBondDenom}
+		// set the reward distribution to be performed during the next block
+		params.BlocksPerDistributionTransmission = int64(1)
+		consumerKeeper.SetParams(bundle.GetCtx(), params)
+
+		// transfer the consumer reward pool to the provider
+		var rewardsPerConsumer sdk.Coin
+
+		// check the consumer pool balance
+		// Note that for a democracy consumer chain the pool may already be filled
+		if pool := bankKeeper.GetAllBalances(
+			bundle.GetCtx(),
+			accountKeeper.GetModuleAccount(bundle.GetCtx(), consumertypes.ConsumerToSendToProviderName).GetAddress(),
+		); pool.Empty() {
+			// if pool is empty, fill it with some tokens
+			rewardsPerConsumer = sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(100))
+			err := bankKeeper.SendCoinsFromAccountToModule(
+				bundle.GetCtx(),
+				bundle.Chain.SenderAccount.GetAddress(),
+				consumertypes.ConsumerToSendToProviderName,
+				sdk.NewCoins(rewardsPerConsumer),
+			)
+			s.Require().NoError(err)
+		} else {
+			// execute the internal rewards distribution
+			// to save the pool's balance before
+			// it gets transferred to the provider in EndBlock
+			consumerKeeper.DistributeRewardsInternally(bundle.GetCtx())
+			pool = bankKeeper.GetAllBalances(
+				bundle.GetCtx(),
+				accountKeeper.GetModuleAccount(bundle.GetCtx(), consumertypes.ConsumerToSendToProviderName).GetAddress(),
+			)
+			s.Require().Len(pool, 1, "consumer reward pool cannot have mutiple token denoms")
+			rewardsPerConsumer = pool[0]
+		}
+
+		// perform the reward transfer
+		bundle.Chain.NextBlock()
+
+		// relay IBC transfer packet from consumer to provider
+		relayAllCommittedPackets(
+			s,
+			bundle.Chain,
+			bundle.TransferPath,
+			transfertypes.PortID,
+			bundle.TransferPath.EndpointA.ChannelID,
+			1,
+		)
+
+		// construct the denom of the reward tokens for the provider
+		prefixedDenom := ibctransfertypes.GetPrefixedDenom(
+			transfertypes.PortID,
+			bundle.TransferPath.EndpointB.ChannelID,
+			rewardsPerConsumer.Denom,
+		)
+		provIBCDenom := ibctransfertypes.ParseDenomTrace(prefixedDenom).IBCDenom()
+
+		// sum the total rewards transferred to the provider
+		totalConsumerRewards = totalConsumerRewards.
+			Add(sdk.NewCoin(provIBCDenom, rewardsPerConsumer.Amount))
+	}
+
+	// Check that the provider receives the rewards of each consumer
+	rewardCoins = providerBankKeeper.GetAllBalances(s.providerCtx(), rewardPool)
+	s.Require().Equal(totalConsumerRewards, rewardCoins, totalConsumerRewards.String(), rewardCoins.String())
+}

tests/integration/setup.go

@@ -136,12 +136,16 @@
 	preProposalKeyAssignment(suite, icstestingutils.FirstConsumerChainID)
 
 	// start consumer chains
-	numConsumers := 5
 	suite.consumerBundles = make(map[string]*icstestingutils.ConsumerBundle)
-	for i := 0; i < numConsumers; i++ {
+	for i := 0; i < icstestingutils.NumConsumers; i++ {
 		bundle := suite.setupConsumerCallback(&suite.Suite, suite.coordinator, i)
 		suite.consumerBundles[bundle.Chain.ChainID] = bundle
 		suite.registerPacketSniffer(bundle.Chain)
+
+		// check that TopN is correctly set for the consumer
+		topN, found := providerKeeper.GetTopN(suite.providerCtx(), bundle.Chain.ChainID)
+		suite.Require().True(found)
+		suite.Require().Equal(bundle.TopN, topN)
 	}
 
 	// initialize each consumer chain with it's corresponding genesis state
@@ -222,7 +226,6 @@
 	)
 	s.Require().True(found, "provider endpoint clientID not found")
 	bundle.Path.EndpointB.ClientID = providerEndpointClientID
-
 	// Set consumer endpoint's clientID
 	consumerKeeper := bundle.GetKeeper()
 	consumerEndpointClientID, found := consumerKeeper.GetProviderClientID(bundle.GetCtx())
@@ -302,34 +305,70 @@
 
 // TODO: Make SetupTransferChannel functional for multiple consumers by pattern matching SetupCCVChannel.
 // See: https://github.com/cosmos/interchain-security/issues/506
+// SetupTransferChannel setup the transfer channel of the first consumer chain among multiple
 func (suite *CCVTestSuite) SetupTransferChannel() {
-	// transfer path will use the same connection as ccv path
+	suite.setupTransferChannel(
+		suite.transferPath,
+		suite.path,
+		suite.consumerApp.GetConsumerKeeper().GetDistributionTransmissionChannel(
+			suite.consumerChain.GetContext(),
+		),
+	)
+}
 
-	suite.transferPath.EndpointA.ClientID = suite.path.EndpointA.ClientID
-	suite.transferPath.EndpointA.ConnectionID = suite.path.EndpointA.ConnectionID
-	suite.transferPath.EndpointB.ClientID = suite.path.EndpointB.ClientID
-	suite.transferPath.EndpointB.ConnectionID = suite.path.EndpointB.ConnectionID
+func (suite *CCVTestSuite) setupTransferChannel(
+	transferPath *ibctesting.Path,
+	ccvPath *ibctesting.Path,
+	channelID string,
+) {
+	// transfer path will use the same connection as ccv path
+	transferPath.EndpointA.ClientID = ccvPath.EndpointA.ClientID
+	transferPath.EndpointA.ConnectionID = ccvPath.EndpointA.ConnectionID
+	transferPath.EndpointB.ClientID = ccvPath.EndpointB.ClientID
+	transferPath.EndpointB.ConnectionID = ccvPath.EndpointB.ConnectionID
 
 	// CCV channel handshake will automatically initiate transfer channel handshake on ACK
 	// so transfer channel will be on stage INIT when CompleteSetupCCVChannel returns.
-	suite.transferPath.EndpointA.ChannelID = suite.consumerApp.GetConsumerKeeper().GetDistributionTransmissionChannel(
-		suite.consumerChain.GetContext())
+	transferPath.EndpointA.ChannelID = channelID
 
 	// Complete TRY, ACK, CONFIRM for transfer path
-	err := suite.transferPath.EndpointB.ChanOpenTry()
+	err := transferPath.EndpointB.ChanOpenTry()
 	suite.Require().NoError(err)
 
-	err = suite.transferPath.EndpointA.ChanOpenAck()
+	err = transferPath.EndpointA.ChanOpenAck()
 	suite.Require().NoError(err)
 
-	err = suite.transferPath.EndpointB.ChanOpenConfirm()
+	err = transferPath.EndpointB.ChanOpenConfirm()
 	suite.Require().NoError(err)
 
 	// ensure counterparty is up to date
-	err = suite.transferPath.EndpointA.UpdateClient()
+	err = transferPath.EndpointA.UpdateClient()
 	suite.Require().NoError(err)
 }
 
+// SetupAllTransferChannel setup all consumer chains transfer channel
+func (suite *CCVTestSuite) SetupAllTransferChannels() {
+	// setup the first consumer transfer channel
+	suite.SetupTransferChannel()
+
+	// setup all the remaining consumers transfer channels
+	for chainID := range suite.consumerBundles {
+		// skip fist consumer
+		if chainID == suite.consumerChain.ChainID {
+			continue
+		}
+
+		// get the bundle for the chain ID
+		bundle := suite.consumerBundles[chainID]
+		// setup the transfer channel
+		suite.setupTransferChannel(
+			bundle.TransferPath,
+			bundle.Path,
+			bundle.App.GetConsumerKeeper().GetDistributionTransmissionChannel(bundle.GetCtx()),
+		)
+	}
+}
+
 func (s CCVTestSuite) validateEndpointsClientConfig(consumerBundle icstestingutils.ConsumerBundle) { //nolint:govet // this is a test so we can copy locks
 	consumerKeeper := consumerBundle.GetKeeper()
 	providerStakingKeeper := s.providerApp.GetTestStakingKeeper()

testutil/ibc_testing/generic_setup.go

@@ -30,10 +30,16 @@ type (
 // and/or democracy consumer app.go implementation. You should not need to modify or replicate this file
 // to run integration tests against your app.go implementations!
 
+const (
+	// Default number of consumer chains
+	NumConsumers = 5
+)
+
 var (
 	FirstConsumerChainID string
 	provChainID          string
 	democConsumerChainID string
+	consumerTopNParams   [NumConsumers]uint32
 )
 
 func init() {
@@ -42,6 +48,9 @@ func init() {
 	FirstConsumerChainID = ibctesting.GetChainID(2)
 	provChainID = ibctesting.GetChainID(1)
 	democConsumerChainID = ibctesting.GetChainID(5000)
+	// TopN parameter values per consumer chain initiated
+	// sorted in ascending order i.e. testchain2, testchain3, ..., testchain6
+	consumerTopNParams = [NumConsumers]uint32{100, 100, 100, 100, 100}
 }
 
 // ConsumerBundle serves as a way to store useful in-mem consumer app chain state
@@ -51,6 +60,7 @@ type ConsumerBundle struct {
 	App          testutil.ConsumerApp
 	Path         *ibctesting.Path
 	TransferPath *ibctesting.Path
+	TopN         uint32
 }
 
 // GetCtx returns the context for the ConsumerBundle
@@ -116,6 +126,9 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 	index int,
 	appIniter ValSetAppIniter,
 ) *ConsumerBundle {
+	// check index isn't bigger that the number of consumers
+	s.Require().LessOrEqual(index, NumConsumers)
+
 	// consumer chain ID
 	chainID := ibctesting.GetChainID(index + 2)
 
@@ -126,6 +139,7 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 
 	prop := testkeeper.GetTestConsumerAdditionProp()
 	prop.ChainId = chainID
+	prop.Top_N = consumerTopNParams[index] // isn't used in CreateConsumerClient
 	// NOTE: the initial height passed to CreateConsumerClient
 	// must be the height on the consumer when InitGenesis is called
 	prop.InitialHeight = clienttypes.Height{RevisionNumber: 0, RevisionHeight: 3}
@@ -135,6 +149,10 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 	)
 	s.Require().NoError(err)
 
+	// set the consumer TopN here since the test suite setup only used the consumer addition prop
+	// to create the consumer genesis, see BeginBlockInit in /x/ccv/provider/keeper/proposal.go.
+	providerKeeper.SetTopN(providerChain.GetContext(), chainID, prop.Top_N)
+
 	// commit the state on the provider chain
 	coordinator.CommitBlock(providerChain)
 
@@ -174,5 +192,6 @@ func AddConsumer[Tp testutil.ProviderApp, Tc testutil.ConsumerApp](
 	return &ConsumerBundle{
 		Chain: testChain,
 		App:   consumerToReturn,
+		TopN:  prop.Top_N,
 	}
 }

testutil/integration/debug_test.go

@@ -299,3 +299,7 @@ func TestTransferConsumerRewardsToDistributionModule(t *testing.T) {
 func TestAllocateTokensToValidator(t *testing.T) {
 	runCCVTestByName(t, "TestAllocateTokensToValidator")
 }
+
+func TestMultiConsumerRewardsDistribution(t *testing.T) {
+	runCCVTestByName(t, "TestMultiConsumerRewardsDistribution")
+}