useAuth: AWS Amplify Auth + React Hooks = Easy Auth Management

This is a short post providing a sample implementation of AWS Amplify authentication management in a React app with hooks.

tl;dr

• When Auth.signIn() succeeds, it sets a cookie with session data that can be accessed by Auth.currentSession(). This doesn't seem to be well documented, but it unlocks the ability to preserve authentication state on browser refresh.
• Raw source code and tests.
• Shout out to useHooks.com for the inspiration on the useAuth hook source code.
• Shout out to Kent C. Dodds for the inspiration on the React hook testing strategy and implementation.

The Problem

The desirable outcome addressed by this article is an auth management strategy that...

1. Centrally manages auth state such that it is easily available to all components.
2. Implements this strategy with React hook syntax.
3. The authentication service is AWS Amplify (AWS Cognito under the hood).
4. Is tested.

One thing I found in my initial time with AWS Amplify is that, upon browser refresh, my app would lose the current authentication state. In short, a logged-in user is logged out on browser refresh. And that is annoying.

Additionally, I couldn't find much written on this issue. It is entirely possible that I missed an important line in the AWS documentation, but the discovering that Auth.currentSession() accessed a session cookie retained in the browser was a major epiphany.

The Hook

// use-auth.js

import React, {
  useState, useEffect, useContext, createContext,
} from 'react';
import { Auth } from '@aws-amplify/auth';

// Implement your particular AWS Amplify configuration
const amplifyConfigurationOptions = {
  userPoolRegion: "REGION",
  userPoolId: "POOL_ID",
  userPoolWebClientId: "CLIENT_ID",
};

Auth.configure(amplifyConfigurationOptions);

const AuthContext = createContext();

// Wrap your app with <ProvideAuth />
export function ProvideAuth({ children }) {
  const auth = useProvideAuth();
  return <AuthContext.Provider value={auth}>{children}</AuthContext.Provider>;
}

// Access auth values and functions with custom useAuth hook
export const useAuth = () => useContext(AuthContext);

function useProvideAuth() {
  const [user, setUser] = useState(null);
  const [isSignedIn, setIsSignedIn] = useState(false);

  useEffect(() => {
    // NOTE: check for user or risk an infinite loop
    if (!user) {
      // On component mount
      // If a session cookie exists
      // Then use it to reset auth state
      Auth.currentSession()
        .then((session) => {
          const {
            idToken,
            accessToken,
          } = session;

          // Define your user schema per your needs
          const user = {
            email: idToken.payload.email,
            username: idToken.payload.preferred_username,
            userId: idToken.payload.sub,
            accessToken: accessToken.jwtToken,
          };

          setIsSignedIn(true);
          setUser(user);
        })
        .catch((err) => {
          // handle it
        });
    }
  }, [user]);

  const signIn = ({ email, password }) => Auth.signIn(email, password)
    .then((cognitoUser) => {
      // Set user data and access token to memory
      const {
        attributes,
        signInUserSession: {
          accessToken,
        },
      } = cognitoUser;

      const user = {
        email: attributes.email,
        username: attributes.preferred_username,
        userId: attributes.sub,
        accessToken: accessToken.jwtToken,
      };

      setIsSignedIn(true);
      setUser(user);

      return user;
    });

  const signOut = () => Auth.signOut()
    .then(() => {
      setIsSignedIn(false);
      setUser(null);
    });

  return {
    user,
    isSignedIn,
    signIn,
    signOut,
  };
}

I am an admitted neophyte when it comes to useEffect, so there may be a better implementation for recovering auth state within this callback. In particular, I initially ran into an infinite loop when calling setUser() because user is one of the callback's dependencies. Happy to hear advice on this one.

The Usage

Much pseudo-code, but you get the idea...

// AppRoot.jsx
import React from 'react';

import App from './app'; // uses <MyComponent />
import { ProvideAuth } from './use-auth';

return (
  <ProvideAuth>
    <App />
  </ProvideAuth>
);

// MyComponent.jsx
import React from 'react';

import { useAuth } from './use-auth';

function MyComponent() {
  const { isSignedIn, user, signIn, signOut } = useAuth();

  return (
    <div>
      <div>{`IsSignedIn: ${isSignedIn}`}</div>
      <div>{`Username: ${user?.username}`}</div>
      {isSignedIn ? (
        <button onClick={signOut} type="button">Sign Out</button>
      ) : (
        <button onClick={signIn} type="button">Sign In</button>
      )}
    </div>
  )
};

The Test

It's perfectly feasible to test a hook in the abstract, but Kent C. Dodds convinced me that it is better to test the hook in its natural habitat... a component.

Essentially, set up an example component that uses the hook, then compose expectations that for the state of that component that could only be achieved by the hook.

// Example Component

import React from 'react';
import { ProvideAuth, useAuth } from '../src/use-auth';

function TestComponent() {
  const {
    user,
    isSignedIn,
    signIn,
    signOut,
  } = useAuth();

  const handleSignIn = () => {
    const mockCreds = {
      email: '[e-mail]',
      password: 'pw',
    }

    signIn(mockCreds);
  }

  const handleSignOut = () => signOut()

  return (
    <div>
      <div>{`IsSignedIn: ${isSignedIn}`}</div>
      <div>{`Username: ${user?.username}`}</div>
      <div>{`AccessToken: ${user?.accessToken}`}</div>
      <button onClick={handleSignIn} type="button">SignInButton</button>
      <button onClick={handleSignOut} type="button">SignOutButton</button>
    </div>
  );
}

function UseAuthExample() {
  return (
    <ProvideAuth>
      <TestComponent />
    </ProvideAuth>
  );
}

export { UseAuthExample };

// use-auth.test.jsx

import React from 'react';
import {
  render, screen, fireEvent, act,
} from '@testing-library/react';
import { Auth } from '@aws-amplify/auth';

import { UseAuthExample } from './UseAuthExample';

describe('useAuth', () => {
  beforeEach(() => {
    jest.clearAllMocks();
  });

  it('should provide default values on load when user is not authenticated', () => {
    const currentSessionMock = jest.fn().mockRejectedValue('No user found.');
    Auth.currentSession = currentSessionMock;

    render(<UseAuthExample />);

    const isSignedIn = screen.getByText(/issignedin/i);
    const username = screen.getByText(/username/i);
    const accessToken = screen.getByText(/accesstoken/i);

    expect(isSignedIn).toHaveTextContent('IsSignedIn: false');
    expect(username).toHaveTextContent('Username:');
    expect(accessToken).toHaveTextContent('AccessToken:');
  });

  it('should provide current user on load when current session is found', async () => {
    const currentSessionMock = jest.fn().mockResolvedValue({
      idToken: {
        payload: {
          email: '[e-mail]',
          preferred_username: 'myuser',
          sub: '1234-abcd',
        },
      },
      accessToken: {
        jwtToken: 'fake-token',
      },
    });
    Auth.currentSession = currentSessionMock;

    await act(async () => {
      render(<UseAuthExample />);
    });

    const isSignedIn = screen.getByText(/issignedin/i);
    const username = screen.getByText(/username/i);
    const accessToken = screen.getByText(/accesstoken/i);

    expect(isSignedIn).toHaveTextContent('IsSignedIn: true');
    expect(username).toHaveTextContent('Username: myuser');
    expect(accessToken).toHaveTextContent('AccessToken: fake-token');
  });

  it('should login the user and update ui', async () => {
    const currentSessionMock = jest.fn().mockRejectedValue('No user found.');
    const signInMock = jest.fn().mockResolvedValue({
      attributes: {
        email: '[e-mail]',
        preferred_username: 'myuser',
        sub: '1234-abcd',
      },
      signInUserSession: {
        accessToken: {
          jwtToken: 'fake-token',
        },
      },
    });
    Auth.currentSession = currentSessionMock;
    Auth.signIn = signInMock;

    render(<UseAuthExample />);

    const isSignedIn = screen.getByText(/issignedin/i);
    const username = screen.getByText(/username/i);
    const accessToken = screen.getByText(/accesstoken/i);

    expect(isSignedIn).toHaveTextContent('IsSignedIn: false');
    expect(username).toHaveTextContent('Username:');
    expect(accessToken).toHaveTextContent('AccessToken:');

    const signInButton = screen.getByText(/signinbutton/i);

    await act(async () => {
      fireEvent.click(signInButton);
    });

    expect(isSignedIn).toHaveTextContent('IsSignedIn: true');
    expect(username).toHaveTextContent('Username: myuser');
    expect(accessToken).toHaveTextContent('AccessToken: fake-token');
  });
});