Fix Elements initialization in React Strict/Concurrent mode

Refs cannot be mutated and used to update state in the same time in rendering phase. As this is side-effect, it can produce various bugs in concurrent mode.

In StrictMode Elements doesn't do transition from null to valid stripe instance. As in StrictMode React renders twice, `final` ref becomes `true`, but `ctx` state isn't changed.

References:
https://reactjs.org/docs/strict-mode.html#detecting-unexpected-side-effects
facebook/react#18003
facebook/react#18545

Author: khmm12

src/components/Elements.test.tsx

@@ -152,6 +152,30 @@ describe('Elements', () => {
     expect(result.current).toBe(mockElements);
   });
 
+  test('allows a transition from null to a valid Stripe object in StrictMode', async () => {
+    let stripeProp: any = null;
+    const spy = jest.fn();
+    const TestComponent = () => (
+      <React.StrictMode>
+        <Elements stripe={stripeProp}>
+          <ElementsConsumer>
+            {({stripe, elements}) => {
+              spy({stripe, elements});
+              return null;
+            }}
+          </ElementsConsumer>
+        </Elements>
+      </React.StrictMode>
+    );
+
+    const {rerender} = render(<TestComponent />);
+    expect(spy).toBeCalledWith({stripe: null, elements: null});
+
+    stripeProp = mockStripe;
+    rerender(<TestComponent />);
+    expect(spy).toBeCalledWith({stripe: mockStripe, elements: mockElements});
+  });
+
   test('works with a Promise resolving to a valid Stripe object', async () => {
     const wrapper = ({children}: any) => (
       <Elements stripe={mockStripePromise}>{children}</Elements>
@@ -225,9 +249,7 @@ describe('Elements', () => {
     expect(mockStripe.elements).toHaveBeenCalledWith({foo: 'foo'});
   });
 
-  // TODO(christopher): support Strict Mode first
-  // eslint-disable-next-line jest/no-disabled-tests
-  test.skip('does not allow updates to options after the Stripe Promise is set in StrictMode', async () => {
+  test('does not allow updates to options after the Stripe Promise is set in StrictMode', async () => {
     // Silence console output so test output is less noisy
     consoleWarn.mockImplementation(() => {});
 

src/components/Elements.tsx

@@ -6,8 +6,8 @@ import React from 'react';
 import PropTypes from 'prop-types';
 
 import {isEqual} from '../utils/isEqual';
-import {usePrevious} from '../utils/usePrevious';
-import {isStripe, isPromise} from '../utils/guards';
+import {usePromiseResolver} from '../utils/usePromiseResolver';
+import {isStripe} from '../utils/guards';
 
 const INVALID_STRIPE_ERROR =
   'Invalid prop `stripe` supplied to `Elements`. We recommend using the `loadStripe` utility from `@stripe/stripe-js`. See https://stripe.com/docs/stripe-js/react#elements-props-stripe for details.';
@@ -23,28 +23,6 @@ const validateStripe = (maybeStripe: unknown): null | stripeJs.Stripe => {
   throw new Error(INVALID_STRIPE_ERROR);
 };
 
-type ParsedStripeProp =
-  | {tag: 'empty'}
-  | {tag: 'sync'; stripe: stripeJs.Stripe}
-  | {tag: 'async'; stripePromise: Promise<stripeJs.Stripe | null>};
-
-const parseStripeProp = (raw: unknown): ParsedStripeProp => {
-  if (isPromise(raw)) {
-    return {
-      tag: 'async',
-      stripePromise: Promise.resolve(raw).then(validateStripe),
-    };
-  }
-
-  const stripe = validateStripe(raw);
-
-  if (stripe === null) {
-    return {tag: 'empty'};
-  }
-
-  return {tag: 'sync', stripe};
-};
-
 interface ElementsContextValue {
   elements: stripeJs.StripeElements | null;
   stripe: stripeJs.Stripe | null;
@@ -66,6 +44,14 @@ export const parseElementsContext = (
   return ctx;
 };
 
+const createElementsContext = (stripe: stripeJs.Stripe | null, options?: stripeJs.StripeElementsOptions) => {
+  const elements = stripe ? stripe.elements(options) : null
+  return {
+    stripe,
+    elements
+  }
+}
+
 interface ElementsProps {
   /**
    * A [Stripe object](https://stripe.com/docs/js/initializing) or a `Promise` resolving to a `Stripe` object.
@@ -99,66 +85,44 @@ interface PrivateElementsProps {
  *
  * @docs https://stripe.com/docs/stripe-js/react#elements-provider
  */
-export const Elements: FunctionComponent<ElementsProps> = ({
-  stripe: rawStripeProp,
-  options,
-  children,
-}: PrivateElementsProps) => {
-  const final = React.useRef(false);
-  const isMounted = React.useRef(true);
-  const parsed = React.useMemo(() => parseStripeProp(rawStripeProp), [
-    rawStripeProp,
-  ]);
-  const [ctx, setContext] = React.useState<ElementsContextValue>(() => ({
-    stripe: null,
-    elements: null,
-  }));
-
-  const prevStripe = usePrevious(rawStripeProp);
-  const prevOptions = usePrevious(options);
-  if (prevStripe !== null) {
-    if (prevStripe !== rawStripeProp) {
+export const Elements: FunctionComponent<ElementsProps> = (props: PrivateElementsProps) => {
+  const { children } = props
+
+  if (props.stripe === undefined) throw new Error(INVALID_STRIPE_ERROR);
+
+  const [inputs, setInputs] = React.useState({ rawStripe: props.stripe, options: props.options })
+  React.useEffect(() => {
+    const { rawStripe, options } = inputs
+    const { stripe: nextRawStripe, options: nextOptions } = props
+
+    const canUpdate = rawStripe === null
+    const hasRawStripeChanged = rawStripe !== nextRawStripe
+    const hasOptionsChanged = !isEqual(options, nextOptions)
+
+    if (hasRawStripeChanged && !canUpdate) {
       console.warn(
         'Unsupported prop change on Elements: You cannot change the `stripe` prop after setting it.'
       );
     }
-    if (!isEqual(options, prevOptions)) {
+
+    if (hasOptionsChanged && !canUpdate) {
       console.warn(
         'Unsupported prop change on Elements: You cannot change the `options` prop after setting the `stripe` prop.'
       );
     }
-  }
 
-  if (!final.current) {
-    if (parsed.tag === 'sync') {
-      final.current = true;
-      setContext({
-        stripe: parsed.stripe,
-        elements: parsed.stripe.elements(options),
-      });
-    }
+    const nextInputs = { rawStripe: nextRawStripe, options: nextOptions }
+    if (hasRawStripeChanged && canUpdate) setInputs(nextInputs)
+  }, [inputs, props])
 
-    if (parsed.tag === 'async') {
-      final.current = true;
-      parsed.stripePromise.then((stripe) => {
-        if (stripe && isMounted.current) {
-          // Only update Elements context if the component is still mounted
-          // and stripe is not null. We allow stripe to be null to make
-          // handling SSR easier.
-          setContext({
-            stripe,
-            elements: stripe.elements(options),
-          });
-        }
-      });
-    }
-  }
+  const [maybeStripe = null] = usePromiseResolver(inputs.rawStripe)
+  const resolvedStripe = validateStripe(maybeStripe)
+  const [ctx, setContext] = React.useState(() => createElementsContext(resolvedStripe, inputs.options));
 
+  const shouldInitialize = resolvedStripe !== null && ctx.stripe === null
   React.useEffect(() => {
-    return (): void => {
-      isMounted.current = false;
-    };
-  }, []);
+    if (shouldInitialize) setContext(createElementsContext(resolvedStripe, inputs.options))
+  }, [shouldInitialize, resolvedStripe, inputs.options])
 
   React.useEffect(() => {
     const anyStripe: any = ctx.stripe;

src/utils/usePrevious.test.tsx

@@ -0,0 +1,60 @@
+import {renderHook, act} from '@testing-library/react-hooks';
+import {usePromiseResolver} from './usePromiseResolver';
+import { mockStripe } from '../../test/mocks';
+
+const createImperativePromise = (): [Promise<unknown>, (value?: unknown) => Promise<void>, (reason?: any) => Promise<void>] => {
+  let resolveFn: (value?: unknown) => Promise<void> = () => Promise.resolve()
+  let rejectFn: (reason?: any) => Promise<void> = () => Promise.resolve()
+
+  const promise = new Promise((resolve, reject) => {
+    const createVoidPromise = () => promise.then(() => undefined, () => undefined)
+
+    resolveFn = (value) => {
+      resolve(value)
+      return createVoidPromise()
+    }
+
+    rejectFn = (reason) => {
+      reject(reason)
+      return createVoidPromise()
+    }
+  })
+
+  return [promise, resolveFn, rejectFn]
+}
+
+describe('usePromiseResolver', () => {
+  let stripe: ReturnType<typeof mockStripe>;
+
+  beforeEach(() => {
+    stripe = mockStripe();
+  })
+
+  it('returns resolved on mount when not promise given', () => {
+    const {result} = renderHook(() => usePromiseResolver(stripe));
+    expect(result.current).toEqual([stripe, undefined, 'resolved'])
+  });
+
+  it('returns pending on mount when promise given', () => {
+    const [promise] = createImperativePromise()
+    const {result} = renderHook(() => usePromiseResolver(promise));
+    expect(result.current).toEqual([undefined, undefined, 'pending'])
+  });
+
+  it('returns resolved when given promise resolved', async () => {
+    const [promise, resolve] = createImperativePromise()
+    const {result} = renderHook(() => usePromiseResolver(promise));
+
+    await act(() => resolve(stripe))
+    expect(result.current).toEqual([stripe, undefined, 'resolved'])
+  });
+
+  it('returns rejected when given promise rejected', async () => {
+    const [promise,, reject] = createImperativePromise()
+    const {result} = renderHook(() => usePromiseResolver(promise));
+
+    const error = new Error('Something went wrong')
+    await act(() => reject(error))
+    expect(result.current).toEqual([undefined, error, 'rejected'])
+  });
+});

src/utils/usePrevious.ts

@@ -0,0 +1,51 @@
+import React from 'react';
+import {isPromise} from '../utils/guards';
+
+type PromisePending = [undefined, undefined, 'pending'];
+type PromiseResolved<T> = [T, undefined, 'resolved'];
+type PromiseRejected = [undefined, any, 'rejected'];
+type PromiseState<T> = PromisePending | PromiseResolved<T> | PromiseRejected;
+
+const createPending = (): PromisePending => [undefined, undefined, 'pending'];
+
+const createResolved = <T>(value: T): PromiseResolved<T> => [
+  value,
+  undefined,
+  'resolved',
+];
+
+const createRejected = (reason: any): PromiseRejected => [
+  undefined,
+  reason,
+  'rejected',
+];
+
+export const usePromiseResolver = <T>(
+  mayBePromise: T | PromiseLike<T>
+): PromiseState<T> => {
+  const [state, setState] = React.useState<PromiseState<T>>(() =>
+    isPromise(mayBePromise) ? createPending() : createResolved(mayBePromise)
+  );
+
+  React.useEffect(() => {
+    if (!isPromise(mayBePromise)) return setState(createResolved(mayBePromise));
+
+    let isMounted = true;
+
+    setState(createPending());
+    mayBePromise
+      .then(
+        (resolved) => createResolved(resolved),
+        (error) => createRejected(error)
+      )
+      .then((nextState) => {
+        if (isMounted) setState(nextState);
+      });
+
+    return () => {
+      isMounted = false;
+    };
+  }, [mayBePromise]);
+
+  return state;
+};