fix(wallet): fetch new keys without re-connecting the wallet (#3467)

apps/trading/components/vega-wallet-connect-button/vega-wallet-connect-button.spec.tsx

@@ -23,26 +23,44 @@ const generateJsx = (context: VegaWalletContextShape) => {
   );
 };
 
-it('Not connected', () => {
-  render(generateJsx({ pubKey: null } as VegaWalletContextShape));
+describe('VegaWalletConnectButton', () => {
+  it('should fire dialog when not connected', () => {
+    render(generateJsx({ pubKey: null } as VegaWalletContextShape));
 
-  const button = screen.getByTestId('connect-vega-wallet');
-  expect(button).toHaveTextContent('Connect Vega wallet');
-  fireEvent.click(button);
-  expect(mockUpdateDialogOpen).toHaveBeenCalled();
-});
+    const button = screen.getByTestId('connect-vega-wallet');
+    expect(button).toHaveTextContent('Connect Vega wallet');
+    fireEvent.click(button);
+    expect(mockUpdateDialogOpen).toHaveBeenCalled();
+  });
 
-it('Connected', async () => {
-  const pubKey = { publicKey: '123456__123456', name: 'test' };
-  render(
-    generateJsx({
+  it('should retrieve keys when connected', async () => {
+    const pubKey = { publicKey: '123456__123456', name: 'test' };
+    const pubKey2 = { publicKey: '123456__123457', name: 'test2' };
+    render(
+      generateJsx({
+        pubKey: pubKey.publicKey,
+        pubKeys: [pubKey],
+        fetchPubKeys: () => Promise.resolve([pubKey, pubKey2]),
+      } as VegaWalletContextShape)
+    );
+
+    const button = screen.getByTestId('manage-vega-wallet');
+    expect(button).toHaveTextContent(truncateByChars(pubKey.publicKey));
+    userEvent.click(button);
+    expect(mockUpdateDialogOpen).not.toHaveBeenCalled();
+  });
+
+  it('should fetch keys when connected', async () => {
+    const pubKey = { publicKey: '123456__123456', name: 'test' };
+    const context = {
       pubKey: pubKey.publicKey,
       pubKeys: [pubKey],
-    } as VegaWalletContextShape)
-  );
+    } as VegaWalletContextShape;
+    render(generateJsx(context));
 
-  const button = screen.getByTestId('manage-vega-wallet');
-  expect(button).toHaveTextContent(truncateByChars(pubKey.publicKey));
-  userEvent.click(button);
-  expect(mockUpdateDialogOpen).not.toHaveBeenCalled();
+    const button = screen.getByTestId('manage-vega-wallet');
+    expect(button).toHaveTextContent(truncateByChars(pubKey.publicKey));
+    userEvent.click(button);
+    expect(mockUpdateDialogOpen).not.toHaveBeenCalled();
+  });
 });

apps/trading/components/vega-wallet-connect-button/vega-wallet-connect-button.tsx

@@ -29,7 +29,7 @@ const MobileWalletButton = ({
   isConnected?: boolean;
   activeKey?: PubKey;
 }) => {
-  const { pubKeys, selectPubKey, disconnect } = useVegaWallet();
+  const { pubKeys, selectPubKey, disconnect, fetchPubKeys } = useVegaWallet();
   const openVegaWalletDialog = useVegaWalletDialogStore(
     (store) => store.openVegaWalletDialog
   );
@@ -46,9 +46,12 @@ const MobileWalletButton = ({
       openVegaWalletDialog();
       setDrawerOpen(false);
     } else {
+      if (fetchPubKeys) {
+        fetchPubKeys();
+      }
       setDrawerOpen(!drawerOpen);
     }
-  }, [drawerOpen, isConnected, openVegaWalletDialog]);
+  }, [drawerOpen, fetchPubKeys, isConnected, openVegaWalletDialog]);
 
   const iconClass = drawerOpen
     ? 'hidden'
@@ -145,8 +148,14 @@ export const VegaWalletConnectButton = () => {
     (store) => store.openVegaWalletDialog
   );
   const openTransferDialog = useTransferDialog((store) => store.open);
-  const { pubKey, pubKeys, selectPubKey, disconnect, isReadOnly } =
-    useVegaWallet();
+  const {
+    pubKey,
+    pubKeys,
+    selectPubKey,
+    disconnect,
+    isReadOnly,
+    fetchPubKeys,
+  } = useVegaWallet();
   const isConnected = pubKey !== null;
 
   const activeKey = useMemo(() => {
@@ -162,7 +171,12 @@ export const VegaWalletConnectButton = () => {
             trigger={
               <DropdownMenuTrigger
                 data-testid="manage-vega-wallet"
-                onClick={() => setDropdownOpen((curr) => !curr)}
+                onClick={() => {
+                  if (fetchPubKeys) {
+                    fetchPubKeys();
+                  }
+                  setDropdownOpen(!dropdownOpen);
+                }}
               >
                 {activeKey && (
                   <span className="uppercase">{activeKey.name}</span>

libs/wallet/src/context.ts

@@ -31,6 +31,9 @@ export interface VegaWalletContextShape {
     /** Transaction payload */
     transaction: Transaction
   ) => Promise<TransactionResponse | null>;
+
+  /** Fetch public keys */
+  fetchPubKeys?: () => Promise<PubKey[] | null>;
 }
 
 export const VegaWalletContext = createContext<

libs/wallet/src/provider.spec.tsx

@@ -50,6 +50,7 @@ describe('VegaWalletProvider', () => {
       connect: expect.any(Function),
       disconnect: expect.any(Function),
       sendTx: expect.any(Function),
+      fetchPubKeys: expect.any(Function || undefined),
     });
 
     // Connect
@@ -77,14 +78,47 @@ describe('VegaWalletProvider', () => {
     expect(spyOnSend).toHaveBeenCalledWith(mockPubKeys[1].publicKey, {});
   });
 
+  it('should fetch new keypairs', async () => {
+    const { result } = setup();
+
+    // Default state
+    expect(result.current).toEqual({
+      pubKey: null,
+      pubKeys: null,
+      isReadOnly: false,
+      selectPubKey: expect.any(Function),
+      connect: expect.any(Function),
+      disconnect: expect.any(Function),
+      sendTx: expect.any(Function),
+      fetchPubKeys: expect.any(Function),
+    });
+
+    // Connect
+    await act(async () => {
+      result.current.connect(restConnector);
+      result.current.selectPubKey(mockPubKeys[0].publicKey);
+    });
+    expect(spyOnConnect).toHaveBeenCalled();
+    expect(result.current.pubKeys).toHaveLength(mockPubKeys.length);
+    expect(result.current.pubKey).toBe(mockPubKeys[0].publicKey);
+
+    // Fetch pub keys
+    mockPubKeys.push({ publicKey: '333', name: 'public key 3' });
+    await act(async () => {
+      result.current.fetchPubKeys && result.current.fetchPubKeys();
+    });
+    expect(result.current.pubKeys).toHaveLength(mockPubKeys.length);
+  });
+
   it('persists selected pubkey and disconnects', async () => {
     const { result } = setup();
     expect(result.current.pubKey).toBe(null);
 
     await act(async () => {
       result.current.connect(restConnector);
+      result.current.selectPubKey(mockPubKeys[0].publicKey);
     });
-    expect(result.current.pubKey).toBe(mockPubKeys[1].publicKey);
+    expect(result.current.pubKey).toBe(mockPubKeys[0].publicKey);
 
     // Disconnect
     await act(async () => {

libs/wallet/src/provider.tsx

@@ -21,7 +21,7 @@ export const VegaWalletProvider = ({ children }: VegaWalletProviderProps) => {
   const [pubKey, setPubKey] = useState<string | null>(null);
   const [isReadOnly, setIsReadOnly] = useState<boolean>(false);
 
-  // Arary of pubkeys retrieved from the connector
+  // Array of public keys retrieved from the connector
   const [pubKeys, setPubKeys] = useState<PubKey[] | null>(null);
 
   // Reference to the current connector instance
@@ -32,6 +32,28 @@ export const VegaWalletProvider = ({ children }: VegaWalletProviderProps) => {
     LocalStorage.setItem(WALLET_KEY, pk);
   }, []);
 
+  const fetchPubKeys = useCallback(async () => {
+    if (!connector.current) {
+      throw new Error('No connector');
+    }
+    try {
+      const keys = await connector.current.connect();
+
+      if (keys?.length) {
+        setPubKeys(keys);
+        setIsReadOnly(connector.current instanceof ViewConnector);
+        return keys;
+      } else {
+        return null;
+      }
+    } catch (err) {
+      if (err instanceof WalletClientError) {
+        throw err;
+      }
+      return null;
+    }
+  }, []);
+
   const connect = useCallback(async (c: VegaConnector) => {
     connector.current = c;
     try {
@@ -95,8 +117,18 @@ export const VegaWalletProvider = ({ children }: VegaWalletProviderProps) => {
       connect,
       disconnect,
       sendTx,
+      fetchPubKeys,
     };
-  }, [isReadOnly, pubKey, pubKeys, selectPubKey, connect, disconnect, sendTx]);
+  }, [
+    isReadOnly,
+    pubKey,
+    pubKeys,
+    selectPubKey,
+    connect,
+    disconnect,
+    sendTx,
+    fetchPubKeys,
+  ]);
 
   return (
     <VegaWalletContext.Provider value={contextValue}>

libs/wallet/src/use-json-rpc-connect.ts

@@ -30,13 +30,13 @@ export const useJsonRpcConnect = (onConnect: () => void) => {
         // Check if wallet is configured for the same chain as the app
         setStatus(Status.GettingChainId);
 
-        // Dont throw in when cypress is running as trading app relies on
+        // Do not throw in when cypress is running as trading app relies on
         // mocks which result in a mismatch between chainId for app and
         // chainId for wallet
         if (!('Cypress' in window)) {
           const chainIdResult = await connector.getChainId();
           if (chainIdResult.chainID !== appChainId) {
-            // Throw wallet error for consitent error handling
+            // Throw wallet error for consistent error handling
             throw ClientErrors.WRONG_NETWORK;
           }
         }