From 8ef512da5239d801fce3ef19d0a055b761601b20 Mon Sep 17 00:00:00 2001
From: im-adithya <imadithyavardhan@gmail.com>
Date: Wed, 25 Sep 2024 09:27:56 +0530
Subject: [PATCH] fix: reload balances after transfer is complete

---
 frontend/src/components/TransferFundsButton.tsx | 6 +++---
 frontend/src/screens/channels/Channels.tsx      | 7 +++++--
 frontend/src/screens/wallet/index.tsx           | 7 +++++--
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/frontend/src/components/TransferFundsButton.tsx b/frontend/src/components/TransferFundsButton.tsx
index 91b24902d..38f44953f 100644
--- a/frontend/src/components/TransferFundsButton.tsx
+++ b/frontend/src/components/TransferFundsButton.tsx
@@ -7,13 +7,13 @@ import { request } from "src/utils/request";
 type TransferFundsButtonProps = {
   channels: Channel[] | undefined;
   albyBalance: AlbyBalance;
-  reloadAlbyBalance: () => void;
+  onTransferComplete: () => Promise<void>;
 } & ButtonProps;
 
 export function TransferFundsButton({
   channels,
   albyBalance,
-  reloadAlbyBalance,
+  onTransferComplete,
   children,
   ...props
 }: TransferFundsButtonProps) {
@@ -46,7 +46,7 @@ export function TransferFundsButton({
               "Content-Type": "application/json",
             },
           });
-          await reloadAlbyBalance();
+          await onTransferComplete();
           toast({
             title:
               "🎉 Funds from Alby shared wallet transferred to your Alby Hub!",
diff --git a/frontend/src/screens/channels/Channels.tsx b/frontend/src/screens/channels/Channels.tsx
index 8c381dbd2..699cb5f7e 100644
--- a/frontend/src/screens/channels/Channels.tsx
+++ b/frontend/src/screens/channels/Channels.tsx
@@ -69,7 +69,7 @@ export default function Channels() {
   useSyncWallet();
   const { data: channels } = useChannels();
   const { data: nodeConnectionInfo } = useNodeConnectionInfo();
-  const { data: balances } = useBalances();
+  const { data: balances, mutate: reloadBalances } = useBalances();
   const { data: albyBalance, mutate: reloadAlbyBalance } = useAlbyBalance();
   const navigate = useNavigate();
   const [nodes, setNodes] = React.useState<Node[]>([]);
@@ -270,7 +270,10 @@ export default function Channels() {
                 variant="outline"
                 channels={channels}
                 albyBalance={albyBalance}
-                reloadAlbyBalance={reloadAlbyBalance}
+                onTransferComplete={async () => {
+                  await reloadAlbyBalance();
+                  await reloadBalances();
+                }}
               >
                 Migrate
               </TransferFundsButton>
diff --git a/frontend/src/screens/wallet/index.tsx b/frontend/src/screens/wallet/index.tsx
index 46d4b83d6..b05511823 100644
--- a/frontend/src/screens/wallet/index.tsx
+++ b/frontend/src/screens/wallet/index.tsx
@@ -26,7 +26,7 @@ import { useInfo } from "src/hooks/useInfo";
 
 function Wallet() {
   const { data: info, hasChannelManagement } = useInfo();
-  const { data: balances } = useBalances();
+  const { data: balances, mutate: reloadBalances } = useBalances();
   const { data: channels } = useChannels();
   const { data: albyBalance, mutate: reloadAlbyBalance } = useAlbyBalance();
 
@@ -61,7 +61,10 @@ function Wallet() {
               <TransferFundsButton
                 channels={channels}
                 albyBalance={albyBalance}
-                reloadAlbyBalance={reloadAlbyBalance}
+                onTransferComplete={async () => {
+                  await reloadAlbyBalance();
+                  await reloadBalances();
+                }}
               >
                 Migrate Funds
               </TransferFundsButton>