From 89e6e2b98a52885df0c9e8769c4fe39477749d01 Mon Sep 17 00:00:00 2001
From: Ajaz Ahmed <53616488+ajansari95@users.noreply.github.com>
Date: Wed, 13 Nov 2024 21:59:43 +0530
Subject: [PATCH] feat: add periodic delegation balance poling icq (#1741)

* add periodic poling icq

* revert redundant callback
---
 x/interchainstaking/keeper/abci.go           | 19 ++++++++++++
 x/interchainstaking/keeper/callbacks_test.go | 31 +++++++++++++++++++-
 x/interchainstaking/keeper/zones.go          | 11 +++++++
 3 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/x/interchainstaking/keeper/abci.go b/x/interchainstaking/keeper/abci.go
index dabc82e68..d99d8f1df 100644
--- a/x/interchainstaking/keeper/abci.go
+++ b/x/interchainstaking/keeper/abci.go
@@ -10,10 +10,12 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/telemetry"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	tmtypes "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint/types"
 
+	"github.com/quicksilver-zone/quicksilver/utils/addressutils"
 	"github.com/quicksilver-zone/quicksilver/x/interchainstaking/types"
 )
 
@@ -47,6 +49,23 @@ func (k *Keeper) BeginBlocker(ctx sdk.Context) {
 			if err := k.GCCompletedUnbondings(ctx, zone); err != nil {
 				k.Logger(ctx).Error("error in GCCompletedUnbondings", "error", err.Error())
 			}
+
+			addressBytes, err := addressutils.AccAddressFromBech32(zone.DelegationAddress.Address, zone.AccountPrefix)
+			if err != nil {
+				k.Logger(ctx).Error("cannot decode bech32 delegation addr", "error", err.Error())
+			}
+			zone.DelegationAddress.IncrementBalanceWaitgroup()
+			k.ICQKeeper.MakeRequest(
+				ctx,
+				zone.ConnectionId,
+				zone.ChainId,
+				types.BankStoreKey,
+				append(banktypes.CreateAccountBalancesPrefix(addressBytes), []byte(zone.BaseDenom)...),
+				sdk.NewInt(-1),
+				types.ModuleName,
+				"accountbalance",
+				0,
+			)
 		}
 
 		connection, found := k.IBCKeeper.ConnectionKeeper.GetConnection(ctx, zone.ConnectionId)
diff --git a/x/interchainstaking/keeper/callbacks_test.go b/x/interchainstaking/keeper/callbacks_test.go
index 7672f0a35..64bdcd56b 100644
--- a/x/interchainstaking/keeper/callbacks_test.go
+++ b/x/interchainstaking/keeper/callbacks_test.go
@@ -1840,7 +1840,7 @@ func TestDepositIntervalCallback(t *testing.T) {
 		err = fmt.Errorf("pagination total exceeds int range: %d", res.Pagination.Total)
 	}
 	suite.NoError(err)
-	suite.Equal(txQueryCount, int(res.Pagination.Total)-3) //nolint:gosec
+	suite.Equal(int(res.Pagination.Total), txQueryCount) //nolint:gosec
 }
 
 func TestDepositIntervalCallbackWithExistingTxs(t *testing.T) {
@@ -2843,6 +2843,35 @@ func (suite *KeeperTestSuite) TestPerfBalanceCallbackUpdate() {
 	})
 }
 
+func (suite *KeeperTestSuite) TestPollDelegationAccountBalanceCallback() {
+	suite.Run("poll delegation account balance", func() {
+		suite.SetupTest()
+		suite.setupTestZones()
+
+		quicksilver := suite.GetQuicksilverApp(suite.chainA)
+		quicksilver.InterchainstakingKeeper.CallbackHandler().RegisterCallbacks()
+		ctx := suite.chainA.GetContext()
+
+		zone, _ := quicksilver.InterchainstakingKeeper.GetZone(ctx, suite.chainB.ChainID)
+		zone.DelegationAddress.Balance = sdk.NewCoins(sdk.NewCoin("uatom", sdkmath.NewInt(500)))
+		quicksilver.InterchainstakingKeeper.SetZone(ctx, &zone)
+		zone.DelegationAddress.IncrementBalanceWaitgroup()
+
+		response := sdk.NewCoin("uatom", sdk.NewInt(101))
+		respbz, err := quicksilver.AppCodec().Marshal(&response)
+		suite.NoError(err)
+
+		address := zone.DelegationAddress.Address
+		accAddr, err := sdk.AccAddressFromBech32(address)
+		suite.NoError(err)
+		data := append(banktypes.CreateAccountBalancesPrefix(accAddr), "uatom"...)
+		quicksilver.InterchainstakingKeeper.SetZone(ctx, &zone)
+
+		err = keeper.AccountBalanceCallback(quicksilver.InterchainstakingKeeper, ctx, respbz, icqtypes.Query{ChainId: suite.chainB.ChainID, Request: data})
+		suite.NoError(err)
+	})
+}
+
 func TestRejectNonADR027(t *testing.T) {
 	registry := codectypes.NewInterfaceRegistry()
 	cdc := codec.NewProtoCodec(registry)
diff --git a/x/interchainstaking/keeper/zones.go b/x/interchainstaking/keeper/zones.go
index 0eee51e18..d26f9ca60 100644
--- a/x/interchainstaking/keeper/zones.go
+++ b/x/interchainstaking/keeper/zones.go
@@ -303,6 +303,17 @@ func (k *Keeper) SetAccountBalanceForDenom(ctx sdk.Context, zone *types.Zone, ad
 			return err
 		}
 		k.Logger(ctx).Info("Matched performance address", "address", address, "wg", zone.PerformanceAddress.BalanceWaitgroup, "balance", zone.PerformanceAddress.Balance)
+	case zone.DelegationAddress != nil && address == zone.DelegationAddress.Address:
+		existing := zone.DelegationAddress.Balance.AmountOf(coin.Denom)
+		err = zone.DelegationAddress.SetBalance(zone.DelegationAddress.Balance.Sub(sdk.NewCoins(sdk.NewCoin(coin.Denom, existing))...).Add(coin)) // reset this denom
+		if err != nil {
+			return err
+		}
+		err = zone.DelegationAddress.DecrementBalanceWaitgroup()
+		if err != nil {
+			return err
+		}
+		k.Logger(ctx).Info("Matched delegation address", "address", address, "wg", zone.DelegationAddress.BalanceWaitgroup, "balance", zone.DelegationAddress.Balance)
 	default:
 		panic("unexpected")
 	}