feat: Implement watchBalance & update useBalance with optional subscribing

Author: wottpal

src/helpers/getBalance.ts

@@ -2,35 +2,77 @@ import { ApiPromise } from '@polkadot/api'
 import { AccountId } from '@polkadot/types/interfaces'
 import { BN, formatBalance } from '@polkadot/util'
 
-/**
- * Returns the native token balance of the given `address`.
- */
-export const getBalance = async (
-  api: ApiPromise,
-  address: string | AccountId | undefined,
-): Promise<{
+export type BalanceData = {
   tokenDecimals: number
   tokenSymbol: string
   freeBalance?: BN
   reservedBalance?: BN
   balance?: BN
   balanceFormatted?: string
-}> => {
-  // Get the token decimals and symbol
-  const tokenDecimals = api.registry.chainDecimals?.[0] || 12
-  const tokenSymbol = api.registry.chainTokens?.[0] || 'Unit'
+}
 
+/**
+ * Returns the native token balance of the given `address`.
+ */
+export const getBalance = async (
+  api: ApiPromise,
+  address: string | AccountId | undefined,
+): Promise<BalanceData> => {
+  const { tokenDecimals, tokenSymbol } = parseBalanceData(api)
   if (!address) {
     return {
       tokenDecimals,
       tokenSymbol,
     }
   }
 
-  // Get the balance
+  // Query the chain and parse data
   const result: any = await api.query.system.account(address)
-  const freeBalance: BN = new BN(result?.data?.free || 0)
-  const reservedBalance: BN = new BN(result?.data?.reserved || 0)
+  const balanceData = parseBalanceData(api, result?.data)
+
+  return balanceData
+}
+
+/**
+ * Watches the native token balance of the given `address` and returns it in a callback.
+ * The returned void function can be used to unsubscribe.
+ */
+export const watchBalance = async (
+  api: ApiPromise,
+  address: string | AccountId | undefined,
+  callback: (data: BalanceData) => void,
+): Promise<VoidFunction | null> => {
+  const { tokenDecimals, tokenSymbol } = parseBalanceData(api)
+  if (!address) {
+    callback({
+      tokenDecimals,
+      tokenSymbol,
+    })
+    return null
+  }
+
+  // Query the chain, parse data, and call the callback
+  const unsubscribe: any = await api.query.system.account(
+    address,
+    ({ data }: any) => {
+      const balanceData = parseBalanceData(api, data)
+      callback(balanceData)
+    },
+  )
+  return unsubscribe
+}
+
+/**
+ * Helper to parse the fetched balance data.
+ */
+const parseBalanceData = (api: ApiPromise, data?: any): BalanceData => {
+  // Get the token decimals and symbol
+  const tokenDecimals = api.registry.chainDecimals?.[0] || 12
+  const tokenSymbol = api.registry.chainTokens?.[0] || 'Unit'
+
+  // Get the balance
+  const freeBalance: BN = new BN(data?.free || 0)
+  const reservedBalance: BN = new BN(data?.reserved || 0)
   const balance = reservedBalance.add(freeBalance)
 
   // Format the balance
@@ -39,11 +81,11 @@ export const getBalance = async (
   })
 
   return {
+    tokenDecimals,
+    tokenSymbol,
     freeBalance,
     reservedBalance,
     balance,
     balanceFormatted,
-    tokenSymbol,
-    tokenDecimals,
   }
 }

src/hooks/useBalance.ts

@@ -1,4 +1,4 @@
-import { getBalance } from '@helpers/getBalance'
+import { BalanceData, getBalance, watchBalance } from '@helpers/getBalance'
 import { AccountId } from '@polkadot/types/interfaces'
 import { BN } from '@polkadot/util'
 import { useInkathon } from '@provider'
@@ -7,39 +7,45 @@ import { useEffect, useState } from 'react'
 /**
  * Hook that returns the native token balance of the given `address`.
  */
-export const useBalance = (address?: string | AccountId) => {
+export const useBalance = (address?: string | AccountId, watch?: boolean) => {
   const { api } = useInkathon()
   const [freeBalance, setFreeBalance] = useState<BN>()
   const [reservedBalance, setReservedBalance] = useState<BN>()
   const [balance, setBalance] = useState<BN>()
   const [balanceFormatted, setBalanceFormatted] = useState<string>()
   const [tokenSymbol, setTokenSymbol] = useState<string>()
   const [tokenDecimals, setTokenDecimals] = useState<number>()
+  const [unsubscribes, setUnsubscribes] = useState<(VoidFunction | null)[]>([])
 
   useEffect(() => {
-    ;(async () => {
-      if (!api) {
-        setFreeBalance(undefined)
-        setReservedBalance(undefined)
-        setBalance(undefined)
-        setBalanceFormatted(undefined)
-        setTokenSymbol(undefined)
-        setTokenDecimals(undefined)
-        return
-      }
-
-      const result = await getBalance(api, address)
-
-      setFreeBalance(result.freeBalance)
-      setReservedBalance(result.reservedBalance)
-      setBalance(result.balance)
+    const updateBalanceData = (data: BalanceData) => {
+      setFreeBalance(data.freeBalance)
+      setReservedBalance(data.reservedBalance)
+      setBalance(data.balance)
       setBalanceFormatted(
-        result.balanceFormatted &&
-          `${result.balanceFormatted} ${result.tokenSymbol}`,
+        data.balanceFormatted && `${data.balanceFormatted} ${data.tokenSymbol}`,
       )
-      setTokenSymbol(result.tokenSymbol)
-      setTokenDecimals(result.tokenDecimals)
-    })()
+      setTokenSymbol(data.tokenSymbol)
+      setTokenDecimals(data.tokenDecimals)
+    }
+
+    if (!api) {
+      updateBalanceData({} as BalanceData)
+      return
+    }
+
+    if (watch) {
+      watchBalance(api, address, updateBalanceData).then((unsubscribe) => {
+        setUnsubscribes((prev) => [...prev, unsubscribe])
+      })
+    } else {
+      getBalance(api, address).then(updateBalanceData)
+    }
+
+    return () => {
+      unsubscribes.forEach((unsubscribe) => unsubscribe?.())
+      setUnsubscribes(() => [])
+    }
   }, [api, address])
 
   return {