From 8500764e730f5f8fa79b9358a17fee8b2df419cd Mon Sep 17 00:00:00 2001
From: juan <juandgoc@gmail.com>
Date: Tue, 26 Jan 2021 14:33:10 -0500
Subject: [PATCH] Added wallet balance to repayinput

---
 src/components/CollateralInput/index.tsx | 37 ++++++++++++++----------
 src/components/RepayInput/index.tsx      |  2 +-
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/src/components/CollateralInput/index.tsx b/src/components/CollateralInput/index.tsx
index a29ed4d..1e288c4 100644
--- a/src/components/CollateralInput/index.tsx
+++ b/src/components/CollateralInput/index.tsx
@@ -1,7 +1,7 @@
 import React, { useEffect, useState } from "react";
 import { cache, ParsedAccount } from "../../contexts/accounts";
 import { useConnectionConfig } from "../../contexts/connection";
-import { useLendingReserves, useUserDeposits } from "../../hooks";
+import {useLendingReserves, useUserBalance, useUserDeposits} from "../../hooks";
 import {
   LendingReserve,
   LendingMarket,
@@ -27,9 +27,11 @@ export default function CollateralInput(props: {
   onLeverage?: (leverage: number) => void;
   onInputChange: (value: number | null) => void;
   hideBalance?: boolean;
+  useWalletBalance?: boolean;
   showLeverageSelector?: boolean;
   leverage?: number;
 }) {
+  const { balance: tokenBalance } = useUserBalance(props.reserve.liquidityMint);
   const { reserveAccounts } = useLendingReserves();
   const { tokenMap } = useConnectionConfig();
   const [collateralReserve, setCollateralReserve] = useState<string>();
@@ -38,21 +40,26 @@ export default function CollateralInput(props: {
   const userDeposits = useUserDeposits();
 
   useEffect(() => {
-    const id: string =
-      cache
-        .byParser(LendingReserveParser)
-        .find((acc) => acc === collateralReserve) || "";
-    const parser = cache.get(id) as ParsedAccount<LendingReserve>;
-    if (parser) {
-      const collateralDeposit = userDeposits.userDeposits.find(
-        (u) =>
-          u.reserve.info.liquidityMint.toBase58() ===
-          parser.info.liquidityMint.toBase58()
-      );
-      if (collateralDeposit) setBalance(collateralDeposit.info.amount);
-      else setBalance(0);
+    if (props.useWalletBalance) {
+      setBalance(tokenBalance)
+    } else {
+      const id: string =
+        cache
+          .byParser(LendingReserveParser)
+          .find((acc) => acc === collateralReserve) || "";
+      const parser = cache.get(id) as ParsedAccount<LendingReserve>;
+
+      if (parser) {
+        const collateralDeposit = userDeposits.userDeposits.find(
+          (u) =>
+            u.reserve.info.liquidityMint.toBase58() ===
+            parser.info.liquidityMint.toBase58()
+        );
+        if (collateralDeposit) setBalance(collateralDeposit.info.amount);
+        else setBalance(0);
+      }
     }
-  }, [collateralReserve, userDeposits]);
+  }, [collateralReserve, userDeposits, tokenBalance, props.useWalletBalance]);
 
   const market = cache.get(props.reserve.lendingMarket) as ParsedAccount<
     LendingMarket
diff --git a/src/components/RepayInput/index.tsx b/src/components/RepayInput/index.tsx
index 6c759b8..7530517 100644
--- a/src/components/RepayInput/index.tsx
+++ b/src/components/RepayInput/index.tsx
@@ -211,7 +211,7 @@ export const RepayInput = (props: {
                 setLastTyped("repay");
               }}
               disabled={true}
-              hideBalance={true}
+              useWalletBalance={true}
             />
           </div>
           <Slider marks={marks} value={pct} onChange={setPct} />