Fix: query finality contract balance

babylonlabs-io · Dec 14, 2024 · b2fdb44 · b2fdb44
1 parent 182df13
commit b2fdb44
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/tests/e2e/bcd_consumer_integration_test.go b/tests/e2e/bcd_consumer_integration_test.go
@@ -297,7 +297,7 @@ func (s *BCDConsumerIntegrationTestSuite) Test6ConsumerFPRewardsGeneration() {
 	s.NotNil(czActivatedBlock)
 
 	// Ensure the staking contract balance is initially empty
-	rewards, err := s.cosmwasmController.QueryStakingContractBalances()
+	rewards, err := s.cosmwasmController.QueryFinalityContractBalances()
 	s.NoError(err)
 	s.Empty(rewards)
 
@@ -344,7 +344,7 @@ func (s *BCDConsumerIntegrationTestSuite) Test6ConsumerFPRewardsGeneration() {
 
 	// Ensure consumer rewards are generated and sent to the staking contract
 	s.Eventually(func() bool {
-		rewards, err := s.cosmwasmController.QueryStakingContractBalances()
+		rewards, err := s.cosmwasmController.QueryFinalityContractBalances()
 		if err != nil {
 			s.T().Logf("failed to query rewards: %s", err.Error())
 			return false

diff --git a/tests/e2e/cosmos-integration-e2e/clientcontroller/cosmwasm/cosmwasm.go b/tests/e2e/cosmos-integration-e2e/clientcontroller/cosmwasm/cosmwasm.go
@@ -584,6 +584,10 @@ func (cc *CosmwasmConsumerController) QueryDelegations() (*ConsumerDelegationsRe
 	return &resp, nil
 }
 
+func (cc *CosmwasmConsumerController) QueryFinalityContractBalances() (sdk.Coins, error) {
+	return cc.QueryBalances(cc.cfg.BtcFinalityContractAddress)
+}
+
 func (cc *CosmwasmConsumerController) QueryStakingContractBalances() (sdk.Coins, error) {
 	return cc.QueryBalances(cc.cfg.BtcStakingContractAddress)
 }