move max button functionality inside SwapAmountInput component

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);
 }

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>

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

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
       />,
     );

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>

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;