From d8a4073228fac3803ea61c552d5ad37b6642a66c Mon Sep 17 00:00:00 2001
From: Sam Ricotta <samanthalricotta@gmail.com>
Date: Wed, 16 Oct 2024 10:33:46 +0300
Subject: [PATCH] fix tests

---
 test/e2e/btc_timestamping_e2e_test.go | 21 ++++++++++++++-------
 x/incentive/keeper/grpc_query_test.go |  9 ++++++---
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/test/e2e/btc_timestamping_e2e_test.go b/test/e2e/btc_timestamping_e2e_test.go
index 2124b8daf..23dee3e07 100644
--- a/test/e2e/btc_timestamping_e2e_test.go
+++ b/test/e2e/btc_timestamping_e2e_test.go
@@ -152,9 +152,9 @@ func (s *BTCTimestampingTestSuite) Test4GenerateAndWithdrawReward() {
 	rgs, err := nonValidatorNode.QueryRewardGauge(submitterReporterAddr)
 	s.NoError(err)
 	submitterRg, reporterRg := rgs[itypes.SubmitterType.String()], rgs[itypes.ReporterType.String()]
-	s.T().Logf("submitter witdhrawable reward: %s, reporter witdhrawable reward: %s before withdrawing", submitterRg.Coins.String(), reporterRg.Coins.String())
-	s.False(submitterRg.Coins.IsZero())
-	s.False(reporterRg.Coins.IsZero())
+	s.T().Logf("submitter witdhrawable reward: %s, reporter witdhrawable reward: %s before withdrawing", convertToRewardGauge(submitterRg).GetWithdrawableCoins().String(), convertToRewardGauge(reporterRg).GetWithdrawableCoins().String())
+	s.False(convertToRewardGauge(submitterRg).IsFullyWithdrawn())
+	s.False(convertToRewardGauge(reporterRg).IsFullyWithdrawn())
 
 	// withdraw submitter reward
 	nonValidatorNode.WithdrawReward(itypes.SubmitterType.String(), initialization.ValidatorWalletName)
@@ -170,8 +170,8 @@ func (s *BTCTimestampingTestSuite) Test4GenerateAndWithdrawReward() {
 	rgs2, err := nonValidatorNode.QueryRewardGauge(submitterReporterAddr)
 	s.NoError(err)
 	submitterRg2 := rgs2[itypes.SubmitterType.String()]
-	s.T().Logf("submitter withdrawable reward: %s after withdrawing", submitterRg2.Coins.String())
-	s.True(rgs2[itypes.SubmitterType.String()].Coins.IsZero())
+	s.T().Logf("submitter withdrawable reward: %s after withdrawing", convertToRewardGauge(submitterRg2).GetWithdrawableCoins().String())
+	s.True(convertToRewardGauge(submitterRg2).IsFullyWithdrawn())
 
 	// withdraw reporter reward
 	nonValidatorNode.WithdrawReward(itypes.ReporterType.String(), initialization.ValidatorWalletName)
@@ -187,8 +187,8 @@ func (s *BTCTimestampingTestSuite) Test4GenerateAndWithdrawReward() {
 	rgs3, err := nonValidatorNode.QueryRewardGauge(submitterReporterAddr)
 	s.NoError(err)
 	reporterRg3 := rgs3[itypes.SubmitterType.String()]
-	s.T().Logf("reporter withdrawable reward: %s after withdrawing", reporterRg3.Coins.String())
-	s.True(rgs3[itypes.ReporterType.String()].Coins.IsZero())
+	s.T().Logf("reporter withdrawable reward: %s after withdrawing", convertToRewardGauge(reporterRg3).GetWithdrawableCoins().String())
+	s.True(convertToRewardGauge(reporterRg3).IsFullyWithdrawn())
 }
 
 func (s *BTCTimestampingTestSuite) Test5Wasm() {
@@ -307,3 +307,10 @@ func (s *BTCTimestampingTestSuite) Test6InterceptFeeCollector() {
 	s.T().Logf("incentive module account's balance after a block: %s", incentiveBalance2.String())
 	s.True(incentiveBalance2.IsAllGTE(incentiveBalance))
 }
+
+func convertToRewardGauge(rg *itypes.RewardGaugesResponse) *itypes.RewardGauge {
+	return &itypes.RewardGauge{
+		Coins:          rg.Coins,
+		WithdrawnCoins: rg.WithdrawnCoins,
+	}
+}
diff --git a/x/incentive/keeper/grpc_query_test.go b/x/incentive/keeper/grpc_query_test.go
index 048a7ad55..8c3d3b4c0 100644
--- a/x/incentive/keeper/grpc_query_test.go
+++ b/x/incentive/keeper/grpc_query_test.go
@@ -20,17 +20,20 @@ func FuzzRewardGaugesQuery(f *testing.F) {
 
 		// generate a list of random RewardGauge map and insert them to KVStore
 		// where in each map, key is stakeholder type and address is the reward gauge
-		rgMaps := []map[string]*types.RewardGauge{}
+		rgMaps := []map[string]*types.RewardGaugesResponse{}
 		sAddrList := []sdk.AccAddress{}
 		numRgMaps := datagen.RandomInt(r, 100)
 		for i := uint64(0); i < numRgMaps; i++ {
-			rgMap := map[string]*types.RewardGauge{}
+			rgMap := map[string]*types.RewardGaugesResponse{}
 			sAddr := datagen.GenRandomAccount().GetAddress()
 			sAddrList = append(sAddrList, sAddr)
 			for i := uint64(0); i <= datagen.RandomInt(r, 4); i++ {
 				sType := datagen.GenRandomStakeholderType(r)
 				rg := datagen.GenRandomRewardGauge(r)
-				rgMap[sType.String()] = rg
+				rgMap[sType.String()] = &types.RewardGaugesResponse{
+					Coins:          rg.Coins,
+					WithdrawnCoins: rg.WithdrawnCoins,
+				}
 
 				keeper.SetRewardGauge(ctx, sType, sAddr, rg)
 			}