feat(hooks): update useAuth to handle more auth hub events

packages/react/src/hooks/__tests__/useAuth.test.ts

@@ -3,13 +3,31 @@ import { Hub } from '@aws-amplify/core';
 import { act, renderHook } from '@testing-library/react-hooks';
 import { useAuth } from '../useAuth';
 
-jest.mock('@aws-amplify/auth');
+const currentAuthenticatedUser = jest.fn();
+
+jest
+  .spyOn(Auth, 'currentAuthenticatedUser')
+  .mockImplementation(currentAuthenticatedUser);
+
+// hub events that return valid user object
+const SUCCESS_EVENTS_WITH_USER = ['signIn', 'signUp', 'autoSignIn'];
+
+// hub events that return error object
+const FAILURE_EVENTS_WITH_ERROR = ['tokenRefresh_failure', 'signIn_failure'];
+
+const mockCognitoUser = {
+  username: 'johndoe',
+  attributes: {
+    phone_number: '+1-234-567-890',
+    email: 'john@doe.com',
+  },
+};
 
 describe('useAuth', () => {
   afterEach(() => jest.clearAllMocks());
 
   it('should return default values when initialized', async () => {
-    (Auth.currentAuthenticatedUser as jest.Mock).mockResolvedValue(undefined);
+    currentAuthenticatedUser.mockResolvedValue(undefined);
 
     const { result, waitForNextUpdate } = renderHook(() => useAuth());
 
@@ -21,23 +39,17 @@ describe('useAuth', () => {
   });
 
   it('should invoke Auth.currentAuthenticatedUser function', async () => {
-    const mockCurrentAuthenticatedUser = jest.fn(() => Promise.resolve());
-
-    (Auth.currentAuthenticatedUser as jest.Mock).mockImplementation(
-      mockCurrentAuthenticatedUser
-    );
+    currentAuthenticatedUser.mockResolvedValue({});
 
     const { waitForNextUpdate } = renderHook(() => useAuth());
 
     await waitForNextUpdate();
 
-    expect(mockCurrentAuthenticatedUser).toHaveBeenCalled();
+    expect(currentAuthenticatedUser).toHaveBeenCalled();
   });
 
   it('should set an error when something unexpected happen', async () => {
-    (Auth.currentAuthenticatedUser as jest.Mock).mockRejectedValue(
-      new Error('Unknown error')
-    );
+    currentAuthenticatedUser.mockRejectedValue(new Error('Unknown error'));
 
     const { result, waitForNextUpdate } = renderHook(() => useAuth());
 
@@ -47,17 +59,7 @@ describe('useAuth', () => {
   });
 
   it('should retrieve a Cognito user', async () => {
-    const mockCognitoUser = {
-      username: 'johndoe',
-      attributes: {
-        phone_number: '+1-234-567-890',
-        email: 'john@doe.com',
-      },
-    };
-
-    (Auth.currentAuthenticatedUser as jest.Mock).mockResolvedValue(
-      mockCognitoUser
-    );
+    currentAuthenticatedUser.mockResolvedValue(mockCognitoUser);
 
     const { result, waitForNextUpdate } = renderHook(() => useAuth());
 
@@ -67,65 +69,93 @@ describe('useAuth', () => {
     expect(result.current.user).toBe(mockCognitoUser);
   });
 
-  it('should receive a Cognito user on Auth.signIn Hub event', async () => {
-    const mockCognitoUser = {
-      username: 'johndoe',
-      attributes: {
-        phone_number: '+1-234-567-890',
-        email: 'john@doe.com',
-      },
-    };
+  it.each(SUCCESS_EVENTS_WITH_USER)(
+    'should receive a Cognito user on %s Hub event',
+    async () => {
+      currentAuthenticatedUser.mockResolvedValue(undefined);
+
+      const { result, waitForNextUpdate } = renderHook(() => useAuth());
+
+      await waitForNextUpdate();
+
+      expect(result.current.user).toBe(undefined);
+
+      // Simulate Auth signIn Hub action
+      act(() => {
+        Hub.dispatch('auth', { event: 'signIn', data: mockCognitoUser });
+      });
 
-    (Auth.currentAuthenticatedUser as jest.Mock).mockResolvedValue(undefined);
+      expect(result.current.user).toBe(mockCognitoUser);
+    }
+  );
+
+  it('should should unset user on Auth.signOut Hub event', async () => {
+    currentAuthenticatedUser.mockResolvedValue(mockCognitoUser);
 
     const { result, waitForNextUpdate } = renderHook(() => useAuth());
 
     await waitForNextUpdate();
 
-    expect(result.current.user).toBe(undefined);
+    expect(result.current.user).toBe(mockCognitoUser);
 
-    // Simulate Auth signIn Hub action
+    // Simulate Auth signOut Hub action
     act(() => {
-      Hub.dispatch(
-        'auth',
-        { event: 'signIn', data: mockCognitoUser },
-        'Auth',
-        Symbol.for('amplify_default')
-      );
+      Hub.dispatch('auth', { event: 'signOut' });
     });
 
-    expect(result.current.user).toBe(mockCognitoUser);
+    expect(result.current.user).toBeUndefined();
   });
 
-  it('should should unset user on Auth.signOut Hub event', async () => {
-    const mockCognitoUser = {
-      username: 'johndoe',
-      attributes: {
-        phone_number: '+1-234-567-890',
-        email: 'john@doe.com',
-      },
-    };
-
-    (Auth.currentAuthenticatedUser as jest.Mock).mockResolvedValue(
-      mockCognitoUser
-    );
-
-    const { result, waitForNextUpdate } = renderHook(() => useAuth());
+  it('invokes Auth.currentAuthenticatedUser on tokenRefresh event', async () => {
+    currentAuthenticatedUser.mockResolvedValue(mockCognitoUser);
 
+    const { waitForNextUpdate } = renderHook(() => useAuth());
     await waitForNextUpdate();
 
-    expect(result.current.user).toBe(mockCognitoUser);
+    // Simulate Auth tokenRefresh Hub action
+    act(() => {
+      Hub.dispatch('auth', { event: 'tokenRefresh' });
+    });
+
+    expect(currentAuthenticatedUser).toHaveBeenCalled();
+  });
+
+  it.each(FAILURE_EVENTS_WITH_ERROR)(
+    'returns error on %s event',
+    async (event) => {
+      currentAuthenticatedUser.mockResolvedValue(mockCognitoUser);
+
+      const { result, waitForNextUpdate } = renderHook(() => useAuth());
+      await waitForNextUpdate();
+
+      act(() => {
+        Hub.dispatch('auth', {
+          event,
+          data: new Error('mock auth error'),
+        });
+      });
+
+      expect(result.current.user).toBeUndefined();
+      expect(result.current.error?.message).toBe('mock auth error');
+    }
+  );
+
+  it('returns error on autoSignIn_failure event', async () => {
+    currentAuthenticatedUser.mockResolvedValue(mockCognitoUser);
+
+    const { result, waitForNextUpdate } = renderHook(() => useAuth());
+    await waitForNextUpdate();
 
-    // Simulate Auth signOut Hub action
     act(() => {
-      Hub.dispatch(
-        'auth',
-        { event: 'signOut' },
-        'Auth',
-        Symbol.for('amplify_default')
-      );
+      // adapted from https://github.com/aws-amplify/amplify-js/blob/272c2c607cc4adb5ddc9421444887bdb382227a0/packages/auth/src/Auth.ts#L274-L278
+      Hub.dispatch('auth', {
+        event: 'autoSignIn_failure',
+        // autoSignIn_failure event only contains `message` but not `payload`.
+        message: 'autoSignInError',
+      });
     });
 
     expect(result.current.user).toBeUndefined();
+    expect(result.current.error?.message).toBe('autoSignInError');
   });
 });

packages/react/src/hooks/useAuth.ts

@@ -1,19 +1,15 @@
 import * as React from 'react';
 
-import Auth, { CognitoUser } from '@aws-amplify/auth';
 import { Hub, HubCallback } from '@aws-amplify/core';
-
-// Exposes relevant CognitoUser properties
-interface AuthUser extends CognitoUser {
-  username: string;
-  attributes: Record<string, string>;
-}
+import { AmplifyUser } from '@aws-amplify/ui';
+import { Auth } from 'aws-amplify';
 
 export interface UseAuthResult {
-  user?: AuthUser;
+  user?: AmplifyUser;
   isLoading: boolean;
   error?: Error;
-  fetch?: () => void;
+  /** @deprecated Fetch is handled automatically, do not use this directly */
+  fetch?: () => Promise<void>;
 }
 
 /**
@@ -27,32 +23,64 @@ export const useAuth = (): UseAuthResult => {
     user: undefined,
   });
 
-  const handleAuth: HubCallback = ({ payload }) => {
-    switch (payload.event) {
-      case 'signIn':
-        return setResult({ user: payload.data as AuthUser, isLoading: false });
-      case 'signOut':
-        return setResult({ isLoading: false });
-      default:
-        break;
+  /**
+   * Hub events like `tokenRefresh` will not give back the user object.
+   * This util will be used to get current user after those events.
+   */
+  const fetchCurrentUser = React.useCallback(async () => {
+    setResult({ user: undefined, isLoading: true, error: undefined });
+
+    try {
+      const user = (await Auth.currentAuthenticatedUser()) as AmplifyUser;
+      setResult({ user, isLoading: false });
+    } catch (e) {
+      const error = e as Error;
+      setResult({ error, isLoading: false });
     }
-  };
+  }, [setResult]);
 
-  const fetch = () => {
-    setResult({ isLoading: true });
+  const handleAuth: HubCallback = React.useCallback(
+    ({ payload }) => {
+      switch (payload.event) {
+        // success events
+        case 'signIn':
+        case 'signUp':
+        case 'autoSignIn':
+          setResult({ user: payload.data as AmplifyUser, isLoading: false });
+          break;
+        case 'signOut':
+          setResult({ user: undefined, isLoading: false });
+          break;
 
-    Auth.currentAuthenticatedUser()
-      .then((user: AuthUser) => setResult({ user, isLoading: false }))
-      .catch((error: Error) => setResult({ error, isLoading: false }));
+        // failure events
+        case 'tokenRefresh_failure':
+        case 'signIn_failure':
+          setResult({ error: payload.data as Error, isLoading: false });
+          break;
+        case 'autoSignIn_failure':
+          // autoSignIn just returns error message. Wrap it to an Error object
+          setResult({ error: new Error(payload.message), isLoading: false });
+          break;
 
-    // Handle Hub Auth events
-    Hub.listen('auth', handleAuth);
+        // events that need another fetch
+        case 'tokenRefresh':
+          fetchCurrentUser();
+          break;
+      }
+    },
+    [setResult, fetchCurrentUser]
+  );
 
-    // Stop listening events on unmount
-    return () => Hub.remove('auth', handleAuth);
-  };
+  React.useEffect(() => {
+    const unsubscribe = Hub.listen('auth', handleAuth, 'useAuth');
+    fetchCurrentUser(); // on init, see if user is already logged in
 
-  React.useEffect(fetch, []);
+    return unsubscribe;
+  }, [handleAuth, fetchCurrentUser]);
 
-  return { ...result, fetch };
+  return {
+    ...result,
+    /** @deprecated Fetch is handled automatically, do not use this directly */
+    fetch: fetchCurrentUser,
+  };
 };