From af68798cd0a51725d2199e63e7c22288c0e37294 Mon Sep 17 00:00:00 2001
From: Alissa Crane <alissacranespam@gmail.com>
Date: Mon, 10 Jun 2024 17:41:36 -0400
Subject: [PATCH] move max button functionality inside SwapAmountInput
 component

---
 site/docs/components/SwapAmountInputContainer.tsx |  7 +------
 site/docs/pages/swap/swap-amount-input.mdx        |  4 +---
 site/docs/pages/swap/types.mdx                    |  1 -
 src/swap/components/SwapAmountInput.test.tsx      |  7 ++-----
 src/swap/components/SwapAmountInput.tsx           | 10 ++++++++--
 src/swap/types.ts                                 |  1 -
 6 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/site/docs/components/SwapAmountInputContainer.tsx b/site/docs/components/SwapAmountInputContainer.tsx
index 939fa1eb2f..e67e07aaba 100644
--- a/site/docs/components/SwapAmountInputContainer.tsx
+++ b/site/docs/components/SwapAmountInputContainer.tsx
@@ -8,7 +8,6 @@ type SwapAmountInputContainer = {
     setAmount: (a: string) => void,
     amount: string,
     tokenBalance: string,
-    onMaxButtonClick: () => void,
   ) => ReactElement;
 };
 
@@ -24,9 +23,5 @@ export default function SwapAmountInputContainer({ children }: SwapAmountInputCo
 
   const tokenBalance = TOKEN_BALANCE_MAP[token?.symbol];
 
-  const onMaxButtonClick = () => {
-    setAmount(tokenBalance);
-  };
-
-  return children(token, setToken, setAmount, amount, tokenBalance, onMaxButtonClick);
+  return children(token, setToken, setAmount, amount, tokenBalance);
 }
diff --git a/site/docs/pages/swap/swap-amount-input.mdx b/site/docs/pages/swap/swap-amount-input.mdx
index 1b43e9699d..125f83af9a 100644
--- a/site/docs/pages/swap/swap-amount-input.mdx
+++ b/site/docs/pages/swap/swap-amount-input.mdx
@@ -31,7 +31,6 @@ The `SwapAmountInput` component is a stylized input field designed for users to
   tokenBalance={tokenBalance}
   onAmountChange={setAmount}
   onTokenSelectorClick={setToken}
-  onMaxButtonClick={onMaxButtonClick}
 />
 ```
 
@@ -74,7 +73,7 @@ The `SwapAmountInput` component is a stylized input field designed for users to
 
 <App>
   <SwapAmountInputContainer>
-    {(token, setToken, setAmount, amount, tokenBalance, onMaxButtonClick) => (
+    {(token, setToken, setAmount, amount, tokenBalance) => (
       <SwapAmountInput
         token={token}
         swappableTokens={[
@@ -110,7 +109,6 @@ The `SwapAmountInput` component is a stylized input field designed for users to
         tokenBalance={tokenBalance}
         onAmountChange={setAmount}
         onTokenSelectorClick={setToken}
-        onMaxButtonClick={onMaxButtonClick}
       />
     )}
   </SwapAmountInputContainer>
diff --git a/site/docs/pages/swap/types.mdx b/site/docs/pages/swap/types.mdx
index 19c6f607f6..13e6b04c36 100644
--- a/site/docs/pages/swap/types.mdx
+++ b/site/docs/pages/swap/types.mdx
@@ -14,7 +14,6 @@ type SwapAmountInputReact = {
   token?: Token; // Selected token
   swappableTokens?: Token[]; // Tokens available for swap
   tokenBalance?: string; // Amount of selected token user owns
-  onMaxButtonClick: () => void; // Callback function when max button is clicked
   onAmountChange: (amount: string) => void; // Callback function when the amount changes
   onTokenSelectorClick: () => void; // Callback function when the token selector is clicked
   disabled?: boolean; // Whether the input is disabled
diff --git a/src/swap/components/SwapAmountInput.test.tsx b/src/swap/components/SwapAmountInput.test.tsx
index 44e1deaa6f..a854d905b2 100644
--- a/src/swap/components/SwapAmountInput.test.tsx
+++ b/src/swap/components/SwapAmountInput.test.tsx
@@ -10,7 +10,6 @@ import { Token } from '../../token';
 
 const setAmountMock = jest.fn();
 const selectTokenClickMock = jest.fn();
-const onMaxButtonClickMock = jest.fn();
 
 const token = {
   address: '0x123' as Address,
@@ -109,7 +108,7 @@ describe('SwapAmountInput Component', () => {
     expect(input.value).toBe('1');
   });
 
-  it('should call onMaxButtonClick when the max button is clicked', async () => {
+  it('should call onAmountChange with tokenBalance when the max button is clicked', async () => {
     render(
       <SwapAmountInput
         token={token}
@@ -118,14 +117,13 @@ describe('SwapAmountInput Component', () => {
         amount="1"
         onAmountChange={setAmountMock}
         onTokenSelectorClick={selectTokenClickMock}
-        onMaxButtonClick={onMaxButtonClickMock}
         tokenBalance="100"
       />,
     );
 
     const maxButton = screen.getByTestId('ockSwapAmountInput_MaxButton');
     fireEvent.click(maxButton);
-    expect(onMaxButtonClickMock).toHaveBeenCalled();
+    expect(setAmountMock).toHaveBeenCalledWith('100');
   });
 
   it('should disable the input when disabled prop is true', async () => {
@@ -137,7 +135,6 @@ describe('SwapAmountInput Component', () => {
         amount="1"
         onAmountChange={setAmountMock}
         onTokenSelectorClick={selectTokenClickMock}
-        onMaxButtonClick={onMaxButtonClickMock}
         disabled
       />,
     );
diff --git a/src/swap/components/SwapAmountInput.tsx b/src/swap/components/SwapAmountInput.tsx
index 6f1d555c58..54dc7c74fa 100644
--- a/src/swap/components/SwapAmountInput.tsx
+++ b/src/swap/components/SwapAmountInput.tsx
@@ -12,7 +12,6 @@ export function SwapAmountInput({
   tokenBalance,
   onAmountChange,
   onTokenSelectorClick,
-  onMaxButtonClick,
   disabled = false,
 }: SwapAmountInputReact) {
   const handleAmountChange = useCallback(
@@ -24,6 +23,12 @@ export function SwapAmountInput({
     [onAmountChange],
   );
 
+  const handleMaxButtonClick = useCallback(() => {
+    if (tokenBalance && isValidAmount(tokenBalance)) {
+      onAmountChange(tokenBalance);
+    }
+  }, [tokenBalance, onAmountChange]);
+
   return (
     <div data-testid="ockSwapAmountInput_Container" className="ock-swapamountinput-container">
       <div className="ock-swapamountinput-row">
@@ -39,7 +44,8 @@ export function SwapAmountInput({
         <button
           data-testid="ockSwapAmountInput_MaxButton"
           className="ock-swapamountinput-maxbutton"
-          onClick={onMaxButtonClick}
+          onClick={handleMaxButtonClick}
+          disabled={tokenBalance === undefined}
         >
           Max
         </button>
diff --git a/src/swap/types.ts b/src/swap/types.ts
index 43f686428f..6e34ab8ba9 100644
--- a/src/swap/types.ts
+++ b/src/swap/types.ts
@@ -89,7 +89,6 @@ export type SwapAmountInputReact = {
   token?: Token;
   swappableTokens: Token[];
   tokenBalance?: string;
-  onMaxButtonClick?: () => void;
   onAmountChange: (amount: string) => void;
   onTokenSelectorClick: () => void;
   disabled?: boolean;