Chore: Wallet Refactor

src/wallet/components/ConnectWallet.test.tsx

@@ -33,8 +33,9 @@ vi.mock('@rainbow-me/rainbowkit', () => ({
   ConnectButton: {
     Custom: ({
       children,
-    }: { children: (props: { openConnectModal: () => void }) => ReactNode }) =>
-      children({ openConnectModal: openConnectModalMock }),
+    }: {
+      children: (props: { openConnectModal: () => void }) => ReactNode;
+    }) => children({ openConnectModal: openConnectModalMock }),
   },
 }));
 
@@ -55,6 +56,7 @@ describe('ConnectWallet', () => {
     });
     vi.mocked(useWalletContext).mockReturnValue({
       isOpen: false,
+      handleClose: vi.fn(),
       setIsOpen: vi.fn(),
     });
     vi.mocked(useOnchainKit).mockReturnValue({
@@ -120,28 +122,46 @@ describe('ConnectWallet', () => {
 
   it('should toggle wallet modal on button click when connected', () => {
     const setIsOpenMock = vi.fn();
+    const handleCloseMock = vi.fn();
     vi.mocked(useAccount).mockReturnValue({
       address: '0x123',
       status: 'connected',
     });
     vi.mocked(useWalletContext).mockReturnValue({
       isOpen: false,
       setIsOpen: setIsOpenMock,
+      handleClose: handleCloseMock,
     });
-    render(
+
+    const { rerender } = render(
       <ConnectWallet text="Connect Wallet">
         <div>Wallet Connected</div>
       </ConnectWallet>,
     );
     const button = screen.getByText('Wallet Connected');
     fireEvent.click(button);
     expect(setIsOpenMock).toHaveBeenCalledWith(true);
+
+    vi.mocked(useWalletContext).mockReturnValue({
+      isOpen: true,
+      setIsOpen: setIsOpenMock,
+      handleClose: handleCloseMock,
+    });
+
+    rerender(
+      <ConnectWallet text="Connect Wallet">
+        <div>Wallet Connected</div>
+      </ConnectWallet>,
+    );
+
+    fireEvent.click(button);
+    expect(handleCloseMock).toHaveBeenCalled();
   });
 
   it('applies ock-bg-secondary-active class when isOpen is true', () => {
     vi.mocked(useWalletContext).mockReturnValue({
       isOpen: true,
-      setIsOpen: vi.fn(),
+      handleClose: vi.fn(),
     });
     vi.mocked(useAccount).mockReturnValue({
       address: '0x123',

src/wallet/components/ConnectWallet.tsx

@@ -30,7 +30,7 @@ export function ConnectWallet({
   const { config = { wallet: { display: undefined } } } = useOnchainKit();
 
   // Core Hooks
-  const { isOpen, setIsOpen } = useWalletContext();
+  const { isOpen, setIsOpen, handleClose } = useWalletContext();
   const [isModalOpen, setIsModalOpen] = useState(false);
   const { address: accountAddress, status } = useAccount();
   const { connectors, connect, status: connectStatus } = useConnect();
@@ -63,8 +63,12 @@ export function ConnectWallet({
 
   // Handles
   const handleToggle = useCallback(() => {
-    setIsOpen(!isOpen);
-  }, [isOpen, setIsOpen]);
+    if (isOpen) {
+      handleClose();
+    } else {
+      setIsOpen(true);
+    }
+  }, [isOpen, handleClose, setIsOpen]);
 
   const handleCloseModal = useCallback(() => {
     setIsModalOpen(false);

src/wallet/components/Wallet.test.tsx

@@ -29,39 +29,25 @@ vi.mock('../../ui/react/internal/hooks/useOutsideClick', () => ({
   useOutsideClick: vi.fn(),
 }));
 
-type useOutsideClickType = ReturnType<
-  typeof vi.fn<
-    (
-      ref: React.RefObject<HTMLElement>,
-      callback: (event: MouseEvent) => void,
-    ) => void
-  >
->;
-
 describe('Wallet Component', () => {
-  let mockSetIsOpen: ReturnType<typeof vi.fn>;
-  let mockOutsideClickCallback: (e: MouseEvent) => void;
+  let mockHandleClose: ReturnType<typeof vi.fn>;
 
   beforeEach(() => {
-    mockSetIsOpen = vi.fn();
+    mockHandleClose = vi.fn();
     (useWalletContext as ReturnType<typeof vi.fn>).mockReturnValue({
       isOpen: true,
-      setIsOpen: mockSetIsOpen,
+      handleClose: mockHandleClose,
+      containerRef: { current: document.createElement('div') },
     });
 
-    (useOutsideClick as unknown as useOutsideClickType).mockImplementation(
-      (_, callback) => {
-        mockOutsideClickCallback = callback;
-      },
-    );
-
     vi.clearAllMocks();
   });
 
   it('should render the Wallet component with ConnectWallet', () => {
     (useWalletContext as ReturnType<typeof vi.fn>).mockReturnValue({
       isOpen: false,
-      setIsOpen: mockSetIsOpen,
+      handleClose: mockHandleClose,
+      containerRef: { current: document.createElement('div') },
     });
 
     render(
@@ -78,6 +64,18 @@ describe('Wallet Component', () => {
   });
 
   it('should close the wallet when clicking outside', () => {
+    const container = document.createElement('div');
+    (useWalletContext as ReturnType<typeof vi.fn>).mockReturnValue({
+      isOpen: true,
+      handleClose: mockHandleClose,
+      containerRef: { current: container },
+    });
+
+    const mockOutsideClickCallback = vi.fn();
+    (useOutsideClick as ReturnType<typeof vi.fn>).mockReturnValue({
+      handleOutsideClick: mockOutsideClickCallback,
+    });
+
     render(
       <Wallet>
         <ConnectWallet />
@@ -91,13 +89,19 @@ describe('Wallet Component', () => {
 
     mockOutsideClickCallback({} as MouseEvent);
 
-    expect(mockSetIsOpen).toHaveBeenCalledWith(false);
+    expect(mockOutsideClickCallback).toHaveBeenCalled();
   });
 
   it('should not trigger click handler when wallet is closed', () => {
     (useWalletContext as ReturnType<typeof vi.fn>).mockReturnValue({
       isOpen: false,
-      setIsOpen: mockSetIsOpen,
+      handleClose: mockHandleClose,
+      containerRef: { current: document.createElement('div') },
+    });
+
+    const mockOutsideClickCallback = vi.fn();
+    (useOutsideClick as ReturnType<typeof vi.fn>).mockReturnValue({
+      handleOutsideClick: mockOutsideClickCallback,
     });
 
     render(
@@ -113,6 +117,6 @@ describe('Wallet Component', () => {
 
     mockOutsideClickCallback({} as MouseEvent);
 
-    expect(mockSetIsOpen).not.toHaveBeenCalled();
+    expect(mockHandleClose).not.toHaveBeenCalled();
   });
 });

src/wallet/components/Wallet.tsx

@@ -1,23 +1,19 @@
+import { useIsMounted } from '@/core-react/internal/hooks/useIsMounted';
+import { useTheme } from '@/core-react/internal/hooks/useTheme';
+import { findComponent } from '@/core-react/internal/utils/findComponent';
+import { cn } from '@/styles/theme';
+import { useOutsideClick } from '@/ui-react/internal/hooks/useOutsideClick';
 import { Children, useMemo, useRef } from 'react';
-import { useIsMounted } from '../../core-react/internal/hooks/useIsMounted';
-import { useTheme } from '../../core-react/internal/hooks/useTheme';
-import { findComponent } from '../../core-react/internal/utils/findComponent';
-import { cn } from '../../styles/theme';
-import { useOutsideClick } from '../../ui/react/internal/hooks/useOutsideClick';
 import type { WalletReact } from '../types';
 import { ConnectWallet } from './ConnectWallet';
 import { WalletDropdown } from './WalletDropdown';
 import { WalletProvider, useWalletContext } from './WalletProvider';
 
-const WalletContent = ({ children, className }: WalletReact) => {
-  const { isOpen, setIsOpen } = useWalletContext();
+function WalletContent({ children, className }: WalletReact) {
+  const { isOpen, handleClose } = useWalletContext();
   const walletContainerRef = useRef<HTMLDivElement>(null);
 
-  useOutsideClick(walletContainerRef, () => {
-    if (isOpen) {
-      setIsOpen(false);
-    }
-  });
+  useOutsideClick(walletContainerRef, handleClose);
 
   const { connect, dropdown } = useMemo(() => {
     const childrenArray = Children.toArray(children);
@@ -36,7 +32,7 @@ const WalletContent = ({ children, className }: WalletReact) => {
       {isOpen && dropdown}
     </div>
   );
-};
+}
 
 export const Wallet = ({ children, className }: WalletReact) => {
   const componentTheme = useTheme();

src/wallet/components/WalletDropdown.test.tsx

@@ -1,8 +1,13 @@
 import '@testing-library/jest-dom';
-import { render, renderHook, screen, waitFor } from '@testing-library/react';
+import {
+  fireEvent,
+  render,
+  renderHook,
+  screen,
+  waitFor,
+} from '@testing-library/react';
 import { type Mock, beforeEach, describe, expect, it, vi } from 'vitest';
 import { useAccount } from 'wagmi';
-
 import {
   IdentityProvider,
   useIdentityContext,
@@ -13,11 +18,17 @@ import { WalletDropdown } from './WalletDropdown';
 import { useWalletContext } from './WalletProvider';
 
 vi.mock('wagmi', () => ({
-  useAccount: vi.fn(),
+  useAccount: vi.fn().mockReturnValue({ address: '0x123' }),
+  WagmiProvider: ({ children }: { children: React.ReactNode }) => (
+    <>{children}</>
+  ),
 }));
 
 vi.mock('./WalletProvider', () => ({
   useWalletContext: vi.fn(),
+  WalletProvider: ({ children }: { children: React.ReactNode }) => (
+    <>{children}</>
+  ),
 }));
 
 vi.mock('../../ui/react/internal/hooks/useBreakpoints', () => ({
@@ -97,4 +108,58 @@ describe('WalletDropdown', () => {
       expect(result.current.address).toEqual(address);
     });
   });
+
+  it('sets animation classes correctly based on isClosing', () => {
+    useWalletContextMock.mockReturnValue({ isOpen: true, isClosing: false });
+    const { rerender } = render(<WalletDropdown>Content</WalletDropdown>);
+    const dropdown = screen.getByTestId('ockWalletDropdown');
+    expect(dropdown).toHaveClass(
+      'fade-in slide-in-from-top-1.5 animate-in duration-300 ease-out',
+    );
+
+    useWalletContextMock.mockReturnValue({ isOpen: true, isClosing: true });
+    rerender(<WalletDropdown>Content</WalletDropdown>);
+    expect(dropdown).toHaveClass(
+      'fade-out slide-out-to-top-1.5 animate-out fill-mode-forwards ease-in-out',
+    );
+  });
+
+  it('should handle wallet closing correctly', async () => {
+    const mockSetIsOpen = vi.fn();
+    const mockSetIsClosing = vi.fn();
+
+    useWalletContextMock.mockReturnValue({
+      isOpen: true,
+      isClosing: false,
+      setIsOpen: mockSetIsOpen,
+      setIsClosing: mockSetIsClosing,
+    });
+
+    const { rerender } = render(
+      <WalletDropdown>
+        <div>Content</div>
+      </WalletDropdown>,
+    );
+
+    const dropdown = screen.getByTestId('ockWalletDropdown');
+    expect(dropdown).toHaveClass('fade-in slide-in-from-top-1.5');
+
+    useWalletContextMock.mockReturnValue({
+      isOpen: true,
+      isClosing: true,
+      setIsOpen: mockSetIsOpen,
+      setIsClosing: mockSetIsClosing,
+    });
+
+    rerender(
+      <WalletDropdown>
+        <div>Content</div>
+      </WalletDropdown>,
+    );
+
+    fireEvent.animationEnd(dropdown);
+
+    expect(mockSetIsOpen).toHaveBeenCalledWith(false);
+    expect(mockSetIsClosing).toHaveBeenCalledWith(false);
+  });
 });

src/wallet/components/WalletDropdown.tsx

@@ -5,10 +5,12 @@ import { Identity } from '../../ui/react/identity/components/Identity';
 import { useBreakpoints } from '../../ui/react/internal/hooks/useBreakpoints';
 import type { WalletDropdownReact } from '../types';
 import { WalletBottomSheet } from './WalletBottomSheet';
+import { useWalletContext } from './WalletProvider';
 
 export function WalletDropdown({ children, className }: WalletDropdownReact) {
   const breakpoint = useBreakpoints();
   const { address } = useAccount();
+  const { isClosing, setIsOpen, setIsClosing } = useWalletContext();
 
   const childrenArray = useMemo(() => {
     return Children.toArray(children).map((child) => {
@@ -39,9 +41,18 @@ export function WalletDropdown({ children, className }: WalletDropdownReact) {
       className={cn(
         pressable.default,
         color.foreground,
-        'absolute right-0 z-10 mt-1 flex w-max min-w-[300px] cursor-default flex-col overflow-hidden rounded-xl',
+        'absolute right-0 z-10 mt-1.5 flex w-max min-w-[300px] cursor-default flex-col overflow-hidden rounded-xl',
+        isClosing
+          ? 'fade-out slide-out-to-top-1.5 animate-out fill-mode-forwards ease-in-out'
+          : 'fade-in slide-in-from-top-1.5 animate-in duration-300 ease-out',
         className,
       )}
+      onAnimationEnd={() => {
+        if (isClosing) {
+          setIsOpen(false);
+          setIsClosing(false);
+        }
+      }}
       data-testid="ockWalletDropdown"
     >
       {childrenArray}

src/wallet/components/WalletDropdownDisconnect.test.tsx

@@ -24,7 +24,9 @@ describe('WalletDropdownDisconnect', () => {
 
   it('renders correctly with default props', () => {
     render(<WalletDropdownDisconnect />);
-    expect(screen.getByText('Disconnect')).toBeInTheDocument();
+    expect(
+      screen.getByText('Disconnect', { selector: 'span' }),
+    ).toBeInTheDocument();
     expect(screen.getByRole('button')).toBeInTheDocument();
   });