[_]: refactor(mail)/simplify mail account key fetching

package.json

@@ -24,7 +24,7 @@
   "dependencies": {
     "@internxt/css-config": "^1.1.0",
     "@internxt/lib": "^1.4.1",
-    "@internxt/sdk": "^1.16.0",
+    "@internxt/sdk": "^1.16.2",
     "@internxt/ui": "^0.1.12",
     "@phosphor-icons/react": "^2.1.10",
     "@reduxjs/toolkit": "^2.11.2",

src/hooks/mail/useMailAccountGuard.test.tsx

@@ -7,7 +7,6 @@ import { createTestStore } from '@/test-utils/createTestStore';
 import { MailService } from '@/services/sdk/mail';
 import { ErrorService } from '@/services/error';
 import { MAIL_NOT_SETUP_CODE } from '@/errors';
-import { getMockedUser } from '@/test-utils/fixtures';
 import { LocalStorageService } from '@/services/local-storage';
 import { MailKeysService } from '@/services/mail-keys';
 import { openEncryptionKeystore } from 'internxt-crypto';
@@ -40,7 +39,8 @@ describe('useMailAccountGuard', () => {
     MailKeysService.instance.clear();
   });
 
-  test('When the user has no email yet, then it should stay in loading state', () => {
+  test('When the keys query is in flight, then it should stay in loading state', () => {
+    vi.spyOn(MailService.instance, 'getMailAccountKeys').mockReturnValue(new Promise(() => undefined));
     const store = createTestStore();
 
     const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
@@ -53,8 +53,7 @@ describe('useMailAccountGuard', () => {
     vi.spyOn(LocalStorageService.instance, 'getMnemonic').mockReturnValue('mnemonic');
     const decryptedKeys = { publicKey: new Uint8Array([1]), secretKey: new Uint8Array([2]) };
     mockedOpenKeystore.mockResolvedValue(decryptedKeys);
-    const user = getMockedUser({ email: 'jane@inxt.me' });
-    const store = createTestStore({ user: { isAuthenticated: true, user } });
+    const store = createTestStore();
 
     const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
 
@@ -66,8 +65,7 @@ describe('useMailAccountGuard', () => {
     vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
     vi.spyOn(LocalStorageService.instance, 'getMnemonic').mockReturnValue('mnemonic');
     mockedOpenKeystore.mockRejectedValue(new Error('bad keystore'));
-    const user = getMockedUser({ email: 'jane@inxt.me' });
-    const store = createTestStore({ user: { isAuthenticated: true, user } });
+    const store = createTestStore();
 
     const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
 
@@ -79,8 +77,7 @@ describe('useMailAccountGuard', () => {
     vi.spyOn(LocalStorageService.instance, 'getMnemonic').mockReturnValue('mnemonic');
     const decryptedKeys = { publicKey: new Uint8Array([1]), secretKey: new Uint8Array([2]) };
     mockedOpenKeystore.mockRejectedValueOnce(new Error('bad keystore')).mockResolvedValueOnce(decryptedKeys);
-    const user = getMockedUser({ email: 'jane@inxt.me' });
-    const store = createTestStore({ user: { isAuthenticated: true, user } });
+    const store = createTestStore();
 
     const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
 
@@ -96,22 +93,47 @@ describe('useMailAccountGuard', () => {
       status: 403,
       code: MAIL_NOT_SETUP_CODE,
     } as never);
-    const user = getMockedUser({ email: 'jane@inxt.me' });
-    const store = createTestStore({ user: { isAuthenticated: true, user } });
+    const store = createTestStore();
 
     const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
 
     await waitFor(() => expect(result.current.status).toBe('not-setup'));
   });
 
+  test('When decrypted keys are already cached for the address, then it should not attempt to load a keystore and should be ready', async () => {
+    vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
+    const getMnemonicSpy = vi.spyOn(LocalStorageService.instance, 'getMnemonic').mockReturnValue('mnemonic');
+    const cachedKeys = { publicKey: new Uint8Array([9]), secretKey: new Uint8Array([8]) };
+    MailKeysService.instance.set(mockKeys.address, cachedKeys);
+    const store = createTestStore();
+
+    const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
+
+    await waitFor(() => expect(result.current.status).toBe('ready'));
+    expect(mockedOpenKeystore).not.toHaveBeenCalled();
+    expect(getMnemonicSpy).not.toHaveBeenCalled();
+    expect(MailKeysService.instance.get(mockKeys.address)).toBe(cachedKeys);
+  });
+
+  test('When the mnemonic is missing, then the status should be error and no keystore should be loaded', async () => {
+    vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
+    vi.spyOn(LocalStorageService.instance, 'getMnemonic').mockReturnValue(null as unknown as string);
+    const store = createTestStore();
+
+    const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
+
+    await waitFor(() => expect(result.current.status).toBe('error'));
+    expect(mockedOpenKeystore).not.toHaveBeenCalled();
+    expect(MailKeysService.instance.get(mockKeys.address)).toBeNull();
+  });
+
   test('When fetching keys fails for another reason, then the status should be error', async () => {
     vi.spyOn(MailService.instance, 'getMailAccountKeys').mockRejectedValue(new Error('Network error'));
     vi.spyOn(ErrorService.instance, 'castError').mockReturnValue({
       message: 'Network error',
       status: 500,
     } as never);
-    const user = getMockedUser({ email: 'jane@inxt.me' });
-    const store = createTestStore({ user: { isAuthenticated: true, user } });
+    const store = createTestStore();
 
     const { result } = renderHook(() => useMailAccountGuard(), { wrapper: createWrapper(store) });
 

src/hooks/mail/useMailAccountGuard.ts

@@ -1,19 +1,14 @@
-import { skipToken } from '@reduxjs/toolkit/query';
 import { useEffect, useRef, useState } from 'react';
 import { KeystoreType, openEncryptionKeystore } from 'internxt-crypto';
 import { useGetMailAccountKeysQuery } from '@/store/api/mail';
-import { useAppSelector } from '@/store/hooks';
 import { MailNotSetupError } from '@/errors';
 import { LocalStorageService } from '@/services/local-storage';
 import { MailKeysService } from '@/services/mail-keys';
 
 export type MailAccountGuardStatus = 'loading' | 'ready' | 'not-setup' | 'error';
 
 export const useMailAccountGuard = (): { status: MailAccountGuardStatus } => {
-  const userEmail = useAppSelector((state) => state.user.user?.email);
-  const { data, error, isLoading, isFetching } = useGetMailAccountKeysQuery(
-    userEmail ? { address: userEmail } : skipToken,
-  );
+  const { data, error, isLoading, isFetching } = useGetMailAccountKeysQuery();
   const lastStartedAddress = useRef<string | null>(null);
   const [isDecrypted, setIsDecrypted] = useState(false);
   const [decryptError, setDecryptError] = useState(false);
@@ -66,7 +61,7 @@ export const useMailAccountGuard = (): { status: MailAccountGuardStatus } => {
     void decrypt();
   }, [address, publicKey, encryptionPrivateKey, decryptError]);
 
-  if (!userEmail || isLoading || isFetching) return { status: 'loading' };
+  if (isLoading || isFetching) return { status: 'loading' };
   if (error instanceof MailNotSetupError) return { status: 'not-setup' };
   if (error || decryptError) return { status: 'error' };
   if (isDecrypted) return { status: 'ready' };

src/hooks/mail/useMailKeys.test.tsx

@@ -0,0 +1,57 @@
+import { renderHook, waitFor } from '@testing-library/react';
+import { describe, test, expect, vi, beforeEach } from 'vitest';
+import { Provider } from 'react-redux';
+import type { PropsWithChildren } from 'react';
+import { useMailKeys } from './useMailKeys';
+import { createTestStore } from '@/test-utils/createTestStore';
+import { MailService } from '@/services/sdk/mail';
+import { MailKeysService } from '@/services/mail-keys';
+
+const createWrapper = (store: ReturnType<typeof createTestStore>) => {
+  return ({ children }: PropsWithChildren) => <Provider store={store}>{children}</Provider>;
+};
+
+const mockKeys = {
+  address: 'jane@inxt.me',
+  publicKey: 'pub',
+  encryptionPrivateKey: 'enc',
+  recoveryPrivateKey: 'rec',
+};
+
+describe('useMailKeys', () => {
+  beforeEach(() => {
+    vi.restoreAllMocks();
+    MailKeysService.instance.clear();
+  });
+
+  test('When the keys query has no data yet, then it should return null', () => {
+    vi.spyOn(MailService.instance, 'getMailAccountKeys').mockReturnValue(new Promise(() => undefined));
+    const store = createTestStore();
+
+    const { result } = renderHook(() => useMailKeys(), { wrapper: createWrapper(store) });
+
+    expect(result.current).toBeNull();
+  });
+
+  test('When the address has decrypted keys cached, then it should return them', async () => {
+    vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
+    const decryptedKeys = { publicKey: new Uint8Array([1]), secretKey: new Uint8Array([2]) };
+    MailKeysService.instance.set(mockKeys.address, decryptedKeys);
+    const store = createTestStore();
+
+    const { result } = renderHook(() => useMailKeys(), { wrapper: createWrapper(store) });
+
+    await waitFor(() => expect(result.current).toBe(decryptedKeys));
+  });
+
+  test('When the address has no cached keys, then it should return null even after data loads', async () => {
+    const fetchSpy = vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
+    const store = createTestStore();
+
+    const { result, rerender } = renderHook(() => useMailKeys(), { wrapper: createWrapper(store) });
+
+    await waitFor(() => expect(fetchSpy).toHaveBeenCalled());
+    rerender();
+    expect(result.current).toBeNull();
+  });
+});

src/hooks/mail/useMailKeys.ts

@@ -1,9 +1,10 @@
 import type { HybridKeyPair } from 'internxt-crypto';
-import { useAppSelector } from '@/store/hooks';
+import { useGetMailAccountKeysQuery } from '@/store/api/mail';
 import { MailKeysService } from '@/services/mail-keys';
 
 export const useMailKeys = (): HybridKeyPair | null => {
-  const userEmail = useAppSelector((state) => state.user.user?.email);
-  if (!userEmail) return null;
-  return MailKeysService.instance.get(userEmail);
+  const { data } = useGetMailAccountKeysQuery();
+  const address = data?.address;
+  if (!address) return null;
+  return MailKeysService.instance.get(address);
 };

src/services/sdk/mail/index.ts

@@ -29,12 +29,13 @@ export class MailService {
   }
 
   /**
-   * Gets the mail account keys for the given address.
+   * Gets the mail account keys. When `address` is omitted the backend returns
+   * the keys for the caller's default address.
    *
-   * @param address - The mail address whose keys should be retrieved.
+   * @param address - Optional. The mail address whose keys should be retrieved.
    * @returns A promise that resolves with the encrypted keys for the address.
    */
-  async getMailAccountKeys(address: string): Promise<MailAccountKeysResponse> {
+  async getMailAccountKeys(address?: string): Promise<MailAccountKeysResponse> {
     return this.client.getMailAccountKeys(address);
   }
 

src/services/sdk/mail/mail.service.test.ts

@@ -224,6 +224,24 @@ describe('Mail Service', () => {
       expect(mockMailClient.getMailAccountKeys).toHaveBeenCalledWith(address);
     });
 
+    test('When fetching keys without an address, then the client should be called without one', async () => {
+      const mockKeys = {
+        address: 'jane@internxt.com',
+        publicKey: 'pub',
+        encryptionPrivateKey: 'enc',
+        recoveryPrivateKey: 'rec',
+      };
+      const mockMailClient = {
+        getMailAccountKeys: vi.fn().mockResolvedValue(mockKeys),
+      } as any;
+      vi.spyOn(SdkManager.instance, 'getMail').mockReturnValue(mockMailClient);
+
+      const result = await MailService.instance.getMailAccountKeys();
+
+      expect(result).toStrictEqual(mockKeys);
+      expect(mockMailClient.getMailAccountKeys).toHaveBeenCalledWith(undefined);
+    });
+
     test('When fetching keys fails, then an error should be thrown', async () => {
       const unexpectedError = new Error('Unexpected error');
       const mockMailClient = {

src/store/api/mail/index.ts

@@ -60,12 +60,12 @@ const patchMailsAfterAction = async ({
 
 export const mailApi = api.injectEndpoints({
   endpoints: (builder) => ({
-    getMailAccountKeys: builder.query<MailAccountKeysResponse, { address: string }>({
-      async queryFn({
-        address,
-      }): Promise<{ data: MailAccountKeysResponse } | { error: MailNotSetupError | FetchMailAccountKeysError }> {
+    getMailAccountKeys: builder.query<MailAccountKeysResponse, { address?: string } | void>({
+      async queryFn(
+        arg,
+      ): Promise<{ data: MailAccountKeysResponse } | { error: MailNotSetupError | FetchMailAccountKeysError }> {
         try {
-          const keys = await MailService.instance.getMailAccountKeys(address);
+          const keys = await MailService.instance.getMailAccountKeys(arg?.address);
           return { data: keys };
         } catch (error) {
           const err = ErrorService.instance.castError(error);

src/store/api/mail/mail.api.test.ts

@@ -467,12 +467,23 @@ describe('Mail API', () => {
       recoveryPrivateKey: 'rec',
     };
 
-    test('When fetching the mail account keys, then it should return the keys', async () => {
-      vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
+    test('When fetching the mail account keys without an address, then it should call the service without one and return the keys', async () => {
+      const spy = vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
+      const store = createTestStore();
+
+      const result = await store.dispatch(mailApi.endpoints.getMailAccountKeys.initiate());
+
+      expect(spy).toHaveBeenCalledWith(undefined);
+      expect(result.data).toStrictEqual(mockKeys);
+    });
+
+    test('When fetching the mail account keys with an explicit address, then it should forward the address to the service', async () => {
+      const spy = vi.spyOn(MailService.instance, 'getMailAccountKeys').mockResolvedValue(mockKeys);
       const store = createTestStore();
 
       const result = await store.dispatch(mailApi.endpoints.getMailAccountKeys.initiate({ address }));
 
+      expect(spy).toHaveBeenCalledWith(address);
       expect(result.data).toStrictEqual(mockKeys);
     });
 
@@ -486,7 +497,7 @@ describe('Mail API', () => {
       } as never);
       const store = createTestStore();
 
-      const result = await store.dispatch(mailApi.endpoints.getMailAccountKeys.initiate({ address }));
+      const result = await store.dispatch(mailApi.endpoints.getMailAccountKeys.initiate());
 
       expect(result.error).toBeInstanceOf(MailNotSetupError);
     });
@@ -500,7 +511,7 @@ describe('Mail API', () => {
       } as never);
       const store = createTestStore();
 
-      const result = await store.dispatch(mailApi.endpoints.getMailAccountKeys.initiate({ address }));
+      const result = await store.dispatch(mailApi.endpoints.getMailAccountKeys.initiate());
 
       expect(result.error).toBeInstanceOf(FetchMailAccountKeysError);
     });