Fix Elements initialization in React Strict/Concurrent mode

src/components/Elements.test.tsx

@@ -53,9 +53,7 @@ describe('Elements', () => {
     expect(mockStripe.elements).toHaveBeenCalledTimes(1);
   });
 
-  // TODO(christopher): support Strict Mode first
-  // eslint-disable-next-line jest/no-disabled-tests
-  test.skip('only creates elements once in Strict Mode', () => {
+  test('only creates elements once in Strict Mode', () => {
     const TestComponent = () => {
       const _ = useElements();
       return <div />;
@@ -83,39 +81,39 @@ describe('Elements', () => {
   });
 
   test('provides elements and stripe with the ElementsConsumer component', () => {
-    expect.assertions(2);
+    const spy = jest.fn()
 
     render(
       <Elements stripe={mockStripe}>
         <ElementsConsumer>
           {(ctx) => {
-            expect(ctx.elements).toBe(mockElements);
-            expect(ctx.stripe).toBe(mockStripe);
-
+            spy(ctx)
             return null;
           }}
         </ElementsConsumer>
       </Elements>
     );
+
+    expect(spy).toBeCalledWith({ stripe: mockStripe, elements: mockElements });
   });
 
   test('provides elements and stripe with the ElementsConsumer component in Strict Mode', () => {
-    expect.assertions(2);
+    const spy = jest.fn()
 
     render(
       <React.StrictMode>
-        <Elements stripe={mockStripe}>
-          <ElementsConsumer>
-            {(ctx) => {
-              expect(ctx.elements).toBe(mockElements);
-              expect(ctx.stripe).toBe(mockStripe);
-
-              return null;
-            }}
-          </ElementsConsumer>
-        </Elements>
+      <Elements stripe={mockStripe}>
+        <ElementsConsumer>
+          {(ctx) => {
+            spy(ctx)
+            return null;
+          }}
+        </ElementsConsumer>
+      </Elements>
       </React.StrictMode>
     );
+
+    expect(spy).toBeCalledWith({ stripe: mockStripe, elements: mockElements });
   });
 
   test('provides given stripe instance on mount', () => {
@@ -152,6 +150,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>
+            {(ctx) => {
+              spy(ctx)
+              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 +247,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;
@@ -99,77 +77,55 @@ 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 stripe = validateStripe(maybeStripe)
+  const [elements, setElements] = React.useState<stripeJs.StripeElements | null>(null);
 
   React.useEffect(() => {
-    return (): void => {
-      isMounted.current = false;
-    };
-  }, []);
+    if (stripe) setElements(stripe.elements(inputs.options))
+  }, [stripe, inputs.options])
 
   React.useEffect(() => {
-    const anyStripe: any = ctx.stripe;
+    const anyStripe: any = stripe;
 
     if (!anyStripe || !anyStripe._registerWrapper) {
       return;
     }
 
     anyStripe._registerWrapper({name: 'react-stripe-js', version: _VERSION});
-  }, [ctx.stripe]);
+  }, [stripe]);
 
+  const ctx: ElementsContextValue = { stripe, elements }
   return (
     <ElementsContext.Provider value={ctx}>{children}</ElementsContext.Provider>
   );

src/utils/usePromiseResolver.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'])
+  });
+});