From 2a88fa1fe5bd7705b8c7e43be7f531edda6cfb67 Mon Sep 17 00:00:00 2001
From: wjrjerome <wjrjerome@babylonchain.io>
Date: Wed, 18 Dec 2024 19:13:23 +1100
Subject: [PATCH] fix: rewards fetch shall fallback to 0 if BBN address never
 had delegations

---
 .../PersonalBalance/PersonalBalance.tsx       | 12 ++----
 .../hooks/client/rpc/queries/useBbnQuery.ts   | 38 +++++++++++++++++--
 src/app/hooks/services/useRewardsService.ts   | 28 --------------
 3 files changed, 39 insertions(+), 39 deletions(-)

diff --git a/src/app/components/PersonalBalance/PersonalBalance.tsx b/src/app/components/PersonalBalance/PersonalBalance.tsx
index 0ba85be2..99a6e679 100644
--- a/src/app/components/PersonalBalance/PersonalBalance.tsx
+++ b/src/app/components/PersonalBalance/PersonalBalance.tsx
@@ -25,12 +25,8 @@ export function PersonalBalance() {
   const {
     balanceQuery: { data: cosmosBalance, isLoading: cosmosBalanceLoading },
   } = useBbnQuery();
-  const { getRewards, claimRewards } = useRewardsService();
-
-  const { data: rewards, isLoading: rewardsLoading } = useQuery({
-    queryKey: [QUERY_KEYS.REWARDS],
-    queryFn: getRewards,
-  });
+  const { claimRewards } = useRewardsService();
+  const { rewardsQuery } = useBbnQuery();
 
   const { data: btcBalance, isLoading: btcBalanceLoading } = useQuery({
     queryKey: [QUERY_KEYS.BTC_BALANCE],
@@ -75,9 +71,9 @@ export function PersonalBalance() {
         <div className="divider mx-0 my-2 md:divider-horizontal" />
 
         <StatItem
-          loading={rewardsLoading}
+          loading={rewardsQuery.isLoading}
           title="BBN Rewards"
-          value={`${ubbnToBbn(rewards ?? 0)} BBN`}
+          value={`${ubbnToBbn(rewardsQuery.data ?? 0)} BBN`}
           actionComponent={{
             title: "Claim",
             onAction: claimRewards,
diff --git a/src/app/hooks/client/rpc/queries/useBbnQuery.ts b/src/app/hooks/client/rpc/queries/useBbnQuery.ts
index 06479c8d..320d4f69 100644
--- a/src/app/hooks/client/rpc/queries/useBbnQuery.ts
+++ b/src/app/hooks/client/rpc/queries/useBbnQuery.ts
@@ -17,6 +17,7 @@ import { useClientQuery } from "../../useClient";
 const BBN_BTCLIGHTCLIENT_TIP_KEY = "BBN_BTCLIGHTCLIENT_TIP";
 const BBN_BALANCE_KEY = "BBN_BALANCE";
 const BBN_REWARDS_KEY = "BBN_REWARDS";
+const REWARD_GAUGE_KEY_BTC_DELEGATION = "btc_delegation";
 
 /**
  * Query service for Babylon which contains all the queries for
@@ -28,7 +29,7 @@ export const useBbnQuery = () => {
 
   /**
    * Gets the rewards from the user's account.
-   * @returns {Promise<Object>} - The rewards from the user's account.
+   * @returns {Promise<number>} - The rewards from the user's account.
    */
   const rewardsQuery = useClientQuery({
     queryKey: [BBN_REWARDS_KEY, bech32Address],
@@ -37,13 +38,44 @@ export const useBbnQuery = () => {
         return undefined;
       }
       const { incentive } = setupIncentiveExtension(queryClient);
-
       const req: incentivequery.QueryRewardGaugesRequest =
         incentivequery.QueryRewardGaugesRequest.fromPartial({
           address: bech32Address,
         });
 
-      return incentive.RewardGauges(req);
+      let rewards;
+      try {
+        rewards = await incentive.RewardGauges(req);
+      } catch (error) {
+        // If error message contains "reward gauge not found", silently return 0
+        // This is to handle the case where the user has no rewards, meaning
+        // they have not staked
+        if (
+          error instanceof Error &&
+          error.message.includes("reward gauge not found")
+        ) {
+          return 0;
+        }
+        throw error;
+      }
+      if (!rewards) {
+        return 0;
+      }
+
+      const coins =
+        rewards.rewardGauges[REWARD_GAUGE_KEY_BTC_DELEGATION]?.coins;
+      if (!coins) {
+        return 0;
+      }
+
+      const withdrawnCoins = rewards.rewardGauges[
+        REWARD_GAUGE_KEY_BTC_DELEGATION
+      ]?.withdrawnCoins.reduce((acc, coin) => acc + Number(coin.amount), 0);
+
+      return (
+        coins.reduce((acc, coin) => acc + Number(coin.amount), 0) -
+        (withdrawnCoins || 0)
+      );
     },
     enabled: Boolean(queryClient && connected && bech32Address),
     staleTime: ONE_MINUTE,
diff --git a/src/app/hooks/services/useRewardsService.ts b/src/app/hooks/services/useRewardsService.ts
index ea32667c..bfdba92f 100644
--- a/src/app/hooks/services/useRewardsService.ts
+++ b/src/app/hooks/services/useRewardsService.ts
@@ -7,39 +7,12 @@ import { BBN_REGISTRY_TYPE_URLS } from "@/utils/wallet/bbnRegistry";
 import { useBbnTransaction } from "../client/rpc/mutation/useBbnTransaction";
 import { useBbnQuery } from "../client/rpc/queries/useBbnQuery";
 
-const REWARD_GAUGE_KEY_BTC_DELEGATION = "btc_delegation";
-
 export const useRewardsService = () => {
   const { bech32Address } = useCosmosWallet();
 
   const { rewardsQuery } = useBbnQuery();
   const { estimateBbnGasFee, sendBbnTx } = useBbnTransaction();
 
-  /**
-   * Gets the rewards from the user's account.
-   * @returns {Promise<number>} The rewards from the user's account.
-   */
-  const getRewards = useCallback(async (): Promise<number> => {
-    const rewards = rewardsQuery.data;
-    if (!rewards) {
-      return 0;
-    }
-
-    const coins = rewards.rewardGauges[REWARD_GAUGE_KEY_BTC_DELEGATION]?.coins;
-    if (!coins) {
-      return 0;
-    }
-
-    const withdrawnCoins = rewards.rewardGauges[
-      REWARD_GAUGE_KEY_BTC_DELEGATION
-    ]?.withdrawnCoins.reduce((acc, coin) => acc + Number(coin.amount), 0);
-
-    return (
-      coins.reduce((acc, coin) => acc + Number(coin.amount), 0) -
-      (withdrawnCoins || 0)
-    );
-  }, [rewardsQuery.data]);
-
   /**
    * Estimates the gas fee for claiming rewards.
    * @returns {Promise<number>} The gas fee for claiming rewards.
@@ -62,7 +35,6 @@ export const useRewardsService = () => {
   }, [bech32Address, sendBbnTx, rewardsQuery]);
 
   return {
-    getRewards,
     claimRewards,
     estimateClaimRewardsGas,
   };