feat: re-type walletCapabilities object (#1238)

coinbase · Sep 16, 2024 · 603d02c · 603d02c
1 parent 7e18c98
commit 603d02c
Show file tree

Hide file tree

Showing 15 changed files with 115 additions and 176 deletions.
diff --git a/.changeset/tough-books-study.md b/.changeset/tough-books-study.md
@@ -0,0 +1,5 @@
+---
+'@coinbase/onchainkit': patch
+---
+
+**feat**: re-typed walletCapabilities object in `OnchainKitConfig`. by @0xAlec #1238
diff --git a/src/OnchainKitConfig.ts b/src/OnchainKitConfig.ts
@@ -9,11 +9,6 @@ export const ONCHAIN_KIT_CONFIG: OnchainKitConfig = {
   chain: baseSepolia,
   rpcUrl: null,
   schemaId: null,
-  walletCapabilities: {
-    hasAtomicBatch: false,
-    hasAuxiliaryFunds: false,
-    hasPaymasterService: false,
-  },
 };
 
 /**

diff --git a/src/OnchainKitProvider.test.tsx b/src/OnchainKitProvider.test.tsx
@@ -1,4 +1,4 @@
-import { act, render, screen, waitFor } from '@testing-library/react';
+import { render, screen, waitFor } from '@testing-library/react';
 import { base } from 'viem/chains';
 import '@testing-library/jest-dom';
 
@@ -11,7 +11,6 @@ import { mock } from 'wagmi/connectors';
 import { setOnchainKitConfig } from './OnchainKitConfig';
 import { OnchainKitProvider } from './OnchainKitProvider';
 import type { EASSchemaUid } from './identity/types';
-import { useCapabilitiesSafe } from './useCapabilitiesSafe';
 import { useOnchainKit } from './useOnchainKit';
 
 const queryClient = new QueryClient();
@@ -26,7 +25,6 @@ const mockConfig = createConfig({
     [base.id]: http(),
   },
 });
-vi.mock('./useCapabilitiesSafe');
 
 const TestComponent = () => {
   const { schemaId, apiKey } = useOnchainKit();
@@ -46,13 +44,9 @@ vi.mock('./OnchainKitConfig', () => ({
     chain: base,
     rpcUrl: null,
     schemaId: null,
-    walletCapabilities: {
-      hasPaymasterService: false,
-      hasAtomicBatch: false,
-      hasAuxiliaryFunds: false,
-    },
   },
 }));
+
 describe('OnchainKitProvider', () => {
   const schemaId: EASSchemaUid = `0x${'1'.repeat(64)}`;
   const apiKey = 'test-api-key';
@@ -116,7 +110,7 @@ describe('OnchainKitProvider', () => {
     }).not.toThrow();
   });
 
-  it('should call setOnchainKitConfig with the correct default values', async () => {
+  it('should call setOnchainKitConfig with the correct values', async () => {
     render(
       <WagmiProvider config={mockConfig}>
         <QueryClientProvider client={queryClient}>
@@ -133,75 +127,6 @@ describe('OnchainKitProvider', () => {
       chain: base,
       rpcUrl: null,
       schemaId,
-      walletCapabilities: {
-        hasPaymasterService: false,
-        hasAtomicBatch: false,
-        hasAuxiliaryFunds: false,
-      },
-    });
-  });
-
-  it('should call setOnchainKitConfig when capabilities are found', async () => {
-    const walletCapabilities = {
-      hasPaymasterService: true,
-      hasAtomicBatch: true,
-      hasAuxiliaryFunds: true,
-    };
-    vi.mocked(useCapabilitiesSafe).mockReturnValue(walletCapabilities);
-    await act(async () => {
-      render(
-        <WagmiProvider config={mockConfig}>
-          <QueryClientProvider client={queryClient}>
-            <OnchainKitProvider
-              chain={base}
-              schemaId={schemaId}
-              apiKey={apiKey}
-            >
-              <TestComponent />
-            </OnchainKitProvider>
-          </QueryClientProvider>
-        </WagmiProvider>,
-      );
-    });
-    expect(setOnchainKitConfig).toHaveBeenCalledWith({
-      address: null,
-      apiKey,
-      chain: base,
-      rpcUrl: null,
-      schemaId,
-      walletCapabilities,
-    });
-  });
-
-  it('should call setOnchainKitConfig when capabilities are not found', async () => {
-    const walletCapabilities = {
-      hasPaymasterService: false,
-      hasAtomicBatch: false,
-      hasAuxiliaryFunds: false,
-    };
-    vi.mocked(useCapabilitiesSafe).mockReturnValue(walletCapabilities);
-    await act(async () => {
-      render(
-        <WagmiProvider config={mockConfig}>
-          <QueryClientProvider client={queryClient}>
-            <OnchainKitProvider
-              chain={base}
-              schemaId={schemaId}
-              apiKey={apiKey}
-            >
-              <TestComponent />
-            </OnchainKitProvider>
-          </QueryClientProvider>
-        </WagmiProvider>,
-      );
-    });
-    expect(setOnchainKitConfig).toHaveBeenCalledWith({
-      address: null,
-      apiKey,
-      chain: base,
-      rpcUrl: null,
-      schemaId,
-      walletCapabilities,
     });
   });
 });
diff --git a/src/OnchainKitProvider.tsx b/src/OnchainKitProvider.tsx
@@ -2,7 +2,6 @@ import { createContext, useMemo } from 'react';
 import { ONCHAIN_KIT_CONFIG, setOnchainKitConfig } from './OnchainKitConfig';
 import { checkHashLength } from './internal/utils/checkHashLength';
 import type { OnchainKitContextType, OnchainKitProviderReact } from './types';
-import { useCapabilitiesSafe } from './useCapabilitiesSafe';
 
 export const OnchainKitContext =
   createContext<OnchainKitContextType>(ONCHAIN_KIT_CONFIG);
@@ -21,7 +20,6 @@ export function OnchainKitProvider({
   if (schemaId && !checkHashLength(schemaId, 64)) {
     throw Error('EAS schemaId must be 64 characters prefixed with "0x"');
   }
-  const walletCapabilities = useCapabilitiesSafe({ chainId: chain.id });
 
   const value = useMemo(() => {
     const onchainKitConfig = {
@@ -30,15 +28,10 @@ export function OnchainKitProvider({
       chain: chain,
       rpcUrl: rpcUrl ?? null,
       schemaId: schemaId ?? null,
-      walletCapabilities: walletCapabilities ?? {
-        hasAtomicBatch: false,
-        hasAuxiliaryFunds: false,
-        hasPaymasterService: false,
-      },
     };
     setOnchainKitConfig(onchainKitConfig);
     return onchainKitConfig;
-  }, [address, apiKey, chain, rpcUrl, schemaId, walletCapabilities]);
+  }, [address, apiKey, chain, rpcUrl, schemaId]);
   return (
     <OnchainKitContext.Provider value={value}>
       {children}

diff --git a/src/constants.ts b/src/constants.ts
@@ -0,0 +1,6 @@
+// Capabilities
+export enum Capabilities {
+  AtomicBatch = 'atomicBatch',
+  AuxiliaryFunds = 'auxiliaryFunds',
+  PaymasterService = 'paymasterService',
+}
diff --git a/src/index.ts b/src/index.ts
@@ -11,5 +11,4 @@ export type {
   OnchainKitConfig,
   OnchainKitContextType,
   OnchainKitProviderReact,
-  WalletCapabilities,
 } from './types';
diff --git a/src/useCapabilitiesSafe.test.ts → ...nternal/hooks/useCapabilitiesSafe.test.ts b/src/useCapabilitiesSafe.test.ts → ...nternal/hooks/useCapabilitiesSafe.test.ts
@@ -15,15 +15,11 @@ vi.mock('wagmi/experimental', () => ({
 describe('useCapabilitiesSafe', () => {
   const mockChainId = 1;
   const walletCapabilitiesTrue = {
-    hasAtomicBatch: true,
-    hasAuxiliaryFunds: true,
-    hasPaymasterService: true,
-  };
-  const walletCapabilitiesFalse = {
-    hasAtomicBatch: false,
-    hasAuxiliaryFunds: false,
-    hasPaymasterService: false,
+    atomicBatch: { supported: true },
+    paymasterService: { supported: true },
+    auxiliaryFunds: { supported: true },
   };
+  const walletCapabilitiesFalse = {};
 
   beforeEach(() => {
     vi.resetAllMocks();

diff --git a/src/internal/hooks/useCapabilitiesSafe.ts b/src/internal/hooks/useCapabilitiesSafe.ts
@@ -0,0 +1,23 @@
+import { useMemo } from 'react';
+import type { WalletCapabilities } from 'viem';
+import { useAccount } from 'wagmi';
+import { useCapabilities } from 'wagmi/experimental';
+import type { UseCapabilitiesSafeParams } from '../../types';
+
+export function useCapabilitiesSafe({
+  chainId,
+}: UseCapabilitiesSafeParams): WalletCapabilities {
+  const { isConnected } = useAccount();
+
+  const { data: capabilities, error } = useCapabilities({
+    query: { enabled: isConnected },
+  });
+
+  return useMemo(() => {
+    if (error || !capabilities || !capabilities[chainId]) {
+      return {};
+    }
+
+    return capabilities[chainId];
+  }, [capabilities, chainId, error]);
+}
diff --git a/src/transaction/components/TransactionProvider.test.tsx b/src/transaction/components/TransactionProvider.test.tsx
@@ -6,7 +6,7 @@ import {
   useWaitForTransactionReceipt,
 } from 'wagmi';
 import { waitForTransactionReceipt } from 'wagmi/actions';
-import { useOnchainKit } from '../../useOnchainKit';
+import { useCapabilitiesSafe } from '../../internal/hooks/useCapabilitiesSafe';
 import { useCallsStatus } from '../hooks/useCallsStatus';
 import { useSendCall } from '../hooks/useSendCall';
 import { useSendCalls } from '../hooks/useSendCalls';
@@ -55,8 +55,8 @@ vi.mock('../hooks/useSendWalletTransactions', () => ({
   useSendWalletTransactions: vi.fn(),
 }));
 
-vi.mock('../../useOnchainKit', () => ({
-  useOnchainKit: vi.fn(),
+vi.mock('../../internal/hooks/useCapabilitiesSafe', () => ({
+  useCapabilitiesSafe: vi.fn(),
 }));
 
 const silenceError = () => {
@@ -156,13 +156,7 @@ describe('TransactionProvider', () => {
     (useWaitForTransactionReceipt as ReturnType<typeof vi.fn>).mockReturnValue({
       receipt: undefined,
     });
-    (useOnchainKit as ReturnType<typeof vi.fn>).mockReturnValue({
-      walletCapabilities: {
-        hasAtomicBatch: false,
-        hasPaymasterService: false,
-        hasAuxiliaryFunds: false,
-      },
-    });
+    (useCapabilitiesSafe as ReturnType<typeof vi.fn>).mockReturnValue({});
   });
 
   it('should emit onError when setLifecycleStatus is called with error', async () => {
@@ -293,10 +287,10 @@ describe('TransactionProvider', () => {
       status: 'pending',
       writeContractsAsync: writeContractsAsyncMock,
     });
-    (useOnchainKit as ReturnType<typeof vi.fn>).mockReturnValue({
-      walletCapabilities: {
-        hasAtomicBatch: true,
-      },
+    (useCapabilitiesSafe as ReturnType<typeof vi.fn>).mockReturnValue({
+      atomicBatch: { supported: true },
+      paymasterService: { supported: true },
+      auxiliaryFunds: { supported: true },
     });
     render(
       <TransactionProvider contracts={[]}>
@@ -336,10 +330,10 @@ describe('TransactionProvider', () => {
 
   it('should update context on handleSubmit', async () => {
     const sendWalletTransactionsMock = vi.fn();
-    (useOnchainKit as ReturnType<typeof vi.fn>).mockReturnValue({
-      walletCapabilities: {
-        hasAtomicBatch: true,
-      },
+    (useCapabilitiesSafe as ReturnType<typeof vi.fn>).mockReturnValue({
+      atomicBatch: { supported: true },
+      paymasterService: { supported: true },
+      auxiliaryFunds: { supported: true },
     });
     (useSendWalletTransactions as ReturnType<typeof vi.fn>).mockReturnValue(
       sendWalletTransactionsMock,
@@ -364,10 +358,10 @@ describe('TransactionProvider', () => {
       status: 'idle',
       writeContractsAsync: writeContractsAsyncMock,
     });
-    (useOnchainKit as ReturnType<typeof vi.fn>).mockReturnValue({
-      walletCapabilities: {
-        hasAtomicBatch: true,
-      },
+    (useCapabilitiesSafe as ReturnType<typeof vi.fn>).mockReturnValue({
+      atomicBatch: { supported: true },
+      paymasterService: { supported: true },
+      auxiliaryFunds: { supported: true },
     });
     render(
       <TransactionProvider contracts={[]}>
@@ -435,10 +429,10 @@ describe('TransactionProvider', () => {
     (useSendWalletTransactions as ReturnType<typeof vi.fn>).mockReturnValue(
       sendWalletTransactionsMock,
     );
-    (useOnchainKit as ReturnType<typeof vi.fn>).mockReturnValue({
-      walletCapabilities: {
-        hasAtomicBatch: true,
-      },
+    (useCapabilitiesSafe as ReturnType<typeof vi.fn>).mockReturnValue({
+      atomicBatch: { supported: true },
+      paymasterService: { supported: true },
+      auxiliaryFunds: { supported: true },
     });
     render(
       <TransactionProvider contracts={[]}>

diff --git a/src/transaction/components/TransactionProvider.tsx b/src/transaction/components/TransactionProvider.tsx
@@ -7,15 +7,17 @@ import {
   useState,
 } from 'react';
 import type { Address } from 'viem';
+import { baseSepolia } from 'viem/chains';
 import {
   useAccount,
   useConfig,
   useSwitchChain,
   useWaitForTransactionReceipt,
 } from 'wagmi';
 import { waitForTransactionReceipt } from 'wagmi/actions';
+import { Capabilities } from '../../constants';
+import { useCapabilitiesSafe } from '../../internal/hooks/useCapabilitiesSafe';
 import { useValue } from '../../internal/hooks/useValue';
-import { useOnchainKit } from '../../useOnchainKit';
 import {
   GENERIC_ERROR_MESSAGE,
   TRANSACTION_TYPE_CALLS,
@@ -77,7 +79,9 @@ export function TransactionProvider({
     : TRANSACTION_TYPE_CONTRACTS;
 
   // Retrieve wallet capabilities
-  const { walletCapabilities } = useOnchainKit();
+  const walletCapabilities = useCapabilitiesSafe({
+    chainId: chainId || baseSepolia.id,
+  }); // defaults to Base Sepolia if not provided
 
   const { switchChainAsync } = useSwitchChain();
 
@@ -127,7 +131,8 @@ export function TransactionProvider({
   // For batched, use statusSendCalls or statusWriteContracts
   // For single, use statusSendCall or statusWriteContract
   const transactionStatus = useMemo(() => {
-    const transactionStatuses = walletCapabilities.hasAtomicBatch
+    const transactionStatuses = walletCapabilities[Capabilities.AtomicBatch]
+      ?.supported
       ? {
           [TRANSACTION_TYPE_CALLS]: statusSendCalls,
           [TRANSACTION_TYPE_CONTRACTS]: statusWriteContracts,
@@ -143,7 +148,7 @@ export function TransactionProvider({
     statusSendCall,
     statusWriteContract,
     transactionType,
-    walletCapabilities.hasAtomicBatch,
+    walletCapabilities[Capabilities.AtomicBatch],
   ]);
 
   // Transaction hash for single transaction (non-batched)