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.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.
@@ -101,74 +87,49 @@ interface PrivateElementsProps {
  */
 export const Elements: FunctionComponent<ElementsProps> = ({
   stripe: rawStripeProp,
-  options,
+  options: optionsProp,
   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) {
+  const [inputs, setInputs] = React.useState({ rawStripe: rawStripeProp, options: optionsProp })
+  const { rawStripe, options } = inputs
+  React.useEffect(() => {
+    const hasRawStripeChanged = rawStripe !== rawStripeProp
+    const hasOptionsChanged = !isEqual(options, optionsProp)
+    const canUpdate = rawStripe === null
+
+    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),
-      });
-    }
+    if (hasRawStripeChanged && canUpdate) setInputs({ rawStripe: rawStripeProp, options: optionsProp })
+  }, [rawStripe, options, rawStripeProp, optionsProp])
 
-    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 = usePromiseResolver(rawStripe)
+  const stripe = validateStripe(maybeStripe)
+  const [ctx, setContext] = React.useState<ElementsContextValue>(() => createElementsContext(null));
 
-  React.useEffect(() => {
-    return (): void => {
-      isMounted.current = false;
-    };
-  }, []);
 
   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]);
+
+  React.useEffect(() => {
+    if (stripe) setContext(createElementsContext(stripe, options))
+  }, [stripe, options])
 
   return (
     <ElementsContext.Provider value={ctx}>{children}</ElementsContext.Provider>

src/utils/usePrevious.test.tsx

@@ -0,0 +1,74 @@
+import {FunctionComponent} from 'react';
+import React from 'react';
+import { act } from 'react-dom/test-utils';
+import {mount, ReactWrapper} from 'enzyme';
+import {usePromiseResolver} from './usePromiseResolver';
+import { mockStripe } from '../../test/mocks';
+
+const TestComponentInner: FunctionComponent<{value: any}> = () => null
+const TestComponent: FunctionComponent<{promiseLike: any}> = ({promiseLike}) => {
+  const value = usePromiseResolver(promiseLike);
+  return <TestComponentInner value={value} />
+};
+
+const getHookValue = (wrapper: ReactWrapper) => wrapper.find(TestComponentInner).prop('value')
+
+const createController = (): [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 value on mount when not promise given', () => {
+    const wrapper = mount(<TestComponent promiseLike={stripe} />);
+    expect(wrapper.find(TestComponentInner).prop('value')).toBe(stripe)
+  });
+
+  it('returns null on mount when promise given', () => {
+    const [promise] = createController()
+    const wrapper = mount(<TestComponent promiseLike={promise} />);
+    expect(getHookValue(wrapper)).toBeNull()
+  });
+
+  it('returns value when given promise resolved', () => {
+    const [promise, resolve] = createController()
+    const wrapper = mount(<TestComponent promiseLike={promise} />);
+
+    return Promise.resolve(act(() => resolve(stripe))).then(() => {
+      wrapper.update()
+      expect(getHookValue(wrapper)).toBe(stripe)
+    })
+  });
+
+  it('returns null when given promise rejected', () => {
+    const [promise,, reject] = createController()
+    const wrapper = mount(<TestComponent promiseLike={promise} />);
+
+    return Promise.resolve(act(() => reject(new Error('Something went wrong')))).then(() => {
+      wrapper.update()
+      expect(getHookValue(wrapper)).toBeNull()
+    })
+  });
+});

src/utils/usePrevious.ts

@@ -0,0 +1,26 @@
+import React from 'react';
+import {isPromise} from '../utils/guards';
+
+export const usePromiseResolver = <T>(mayBePromise: T | PromiseLike<T>): T | null => {
+  const [resolved, setResolved] = React.useState(() => {
+    return isPromise(mayBePromise) ? null : mayBePromise
+  })
+
+  React.useEffect(() => {
+    if (!isPromise(mayBePromise)) return setResolved(mayBePromise)
+
+    let isMounted = true
+
+    setResolved(null)
+    mayBePromise.then(resolvedValue => {
+      if (isMounted) setResolved(resolvedValue)
+    }, () => undefined)
+
+    return () => {
+      isMounted = false
+    }
+
+  }, [mayBePromise])
+
+  return resolved
+};