From 65b05dd58bd98650fd2c08e17ef8a98a8cbcf560 Mon Sep 17 00:00:00 2001
From: Christopher Henn <56489932+christopher-stripe@users.noreply.github.com>
Date: Thu, 28 Jan 2021 15:12:51 -0800
Subject: [PATCH] Fix Element effect cleanup when using async stripe prop
 (#153)

---
 src/components/createElementComponent.test.tsx | 17 ++++++++++++++++-
 src/components/createElementComponent.tsx      |  6 ++----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/components/createElementComponent.test.tsx b/src/components/createElementComponent.test.tsx
index 726e18c..ba82573 100644
--- a/src/components/createElementComponent.test.tsx
+++ b/src/components/createElementComponent.test.tsx
@@ -1,5 +1,5 @@
 import React from 'react';
-import {render} from '@testing-library/react';
+import {render, act} from '@testing-library/react';
 
 import {Elements} from './Elements';
 import createElementComponent from './createElementComponent';
@@ -405,5 +405,20 @@ describe('createElementComponent', () => {
       unmount2();
       expect(mockElement.destroy).toHaveBeenCalled();
     });
+
+    it('destroys an existing Element when the component unmounts with an async stripe prop', async () => {
+      const stripePromise = Promise.resolve(mockStripe);
+
+      const {unmount} = render(
+        <Elements stripe={stripePromise}>
+          <CardElement />
+        </Elements>
+      );
+
+      await act(() => stripePromise);
+
+      unmount();
+      expect(mockElement.destroy).toHaveBeenCalled();
+    });
   });
 });
diff --git a/src/components/createElementComponent.tsx b/src/components/createElementComponent.tsx
index 76ce7f9..36a0ca2 100644
--- a/src/components/createElementComponent.tsx
+++ b/src/components/createElementComponent.tsx
@@ -115,11 +115,9 @@ const createElementComponent = (
     }, [options]);
 
     React.useEffect(() => {
-      const element = elementRef.current;
-
       return () => {
-        if (element) {
-          element.destroy();
+        if (elementRef.current) {
+          elementRef.current.destroy();
         }
       };
     }, []);