From 4a5f45a7a7540e1a330019683d293ab00bce5cc5 Mon Sep 17 00:00:00 2001
From: Dusty Greif <dgreif@users.noreply.github.com>
Date: Thu, 8 Apr 2021 12:43:29 -0700
Subject: [PATCH] Check default portal still exists before using it (#1153)

---
 src/Portal/Portal.tsx           | 18 +++++++++---------
 src/__tests__/Portal.tsx        |  4 ----
 src/stories/Overlay.stories.tsx |  4 ----
 src/stories/Portal.stories.tsx  |  3 ---
 4 files changed, 9 insertions(+), 20 deletions(-)

diff --git a/src/Portal/Portal.tsx b/src/Portal/Portal.tsx
index 6ea3c405395..567551ce0ee 100644
--- a/src/Portal/Portal.tsx
+++ b/src/Portal/Portal.tsx
@@ -12,19 +12,18 @@ const portalRootRegistry: {[key: string]: Element} = {}
  * @param name The name of the container, to be used with the `containerName` prop on the Portal Component.
  * If name is not specified, registers the default portal root.
  */
-export function registerPortalRoot(root: Element | undefined, name?: string): void {
-  if (root instanceof Element) {
-    portalRootRegistry[name ?? DEFAULT_PORTAL_CONTAINER_NAME] = root
-  } else {
-    delete portalRootRegistry[name ?? DEFAULT_PORTAL_CONTAINER_NAME]
-  }
+export function registerPortalRoot(root: Element, name = DEFAULT_PORTAL_CONTAINER_NAME): void {
+  portalRootRegistry[name] = root
 }
 
 // Ensures that a default portal root exists and is registered. If a DOM element exists
-// with id __primerPortalRoot__, allow that element to serve as the default portl root.
+// with id __primerPortalRoot__, allow that element to serve as the default portal root.
 // Otherwise, create that element and attach it to the end of document.body.
 function ensureDefaultPortal() {
-  if (!(DEFAULT_PORTAL_CONTAINER_NAME in portalRootRegistry)) {
+  if (
+    !(DEFAULT_PORTAL_CONTAINER_NAME in portalRootRegistry) ||
+    !document.body.contains(portalRootRegistry[DEFAULT_PORTAL_CONTAINER_NAME])
+  ) {
     let defaultPortalContainer = document.getElementById(PRIMER_PORTAL_ROOT_ID)
     if (!(defaultPortalContainer instanceof Element)) {
       defaultPortalContainer = document.createElement('div')
@@ -36,7 +35,8 @@ function ensureDefaultPortal() {
         document.body.appendChild(defaultPortalContainer)
       }
     }
-    portalRootRegistry[DEFAULT_PORTAL_CONTAINER_NAME] = defaultPortalContainer
+
+    registerPortalRoot(defaultPortalContainer)
   }
 }
 
diff --git a/src/__tests__/Portal.tsx b/src/__tests__/Portal.tsx
index fa3c9726293..6ff4de14d81 100644
--- a/src/__tests__/Portal.tsx
+++ b/src/__tests__/Portal.tsx
@@ -5,10 +5,6 @@ import React from 'react'
 import BaseStyles from '../BaseStyles'
 
 describe('Portal', () => {
-  afterEach(() => {
-    // since the registry is global, reset after each test
-    registerPortalRoot(undefined)
-  })
   it('renders a default portal into document.body (no BaseStyles present)', () => {
     const {baseElement} = render(<Portal>123test123</Portal>)
     const generatedRoot = baseElement.querySelector('#__primerPortalRoot__')
diff --git a/src/stories/Overlay.stories.tsx b/src/stories/Overlay.stories.tsx
index a686f66caf6..e2dbcd31640 100644
--- a/src/stories/Overlay.stories.tsx
+++ b/src/stories/Overlay.stories.tsx
@@ -3,7 +3,6 @@ import React, {useState, useRef} from 'react'
 import {Meta} from '@storybook/react'
 import styled from 'styled-components'
 
-import {registerPortalRoot} from '../Portal'
 import {BaseStyles, Overlay, Button, Text, ButtonDanger, ThemeProvider, Position, Flex} from '..'
 
 export default {
@@ -11,9 +10,6 @@ export default {
   component: Overlay,
   decorators: [
     Story => {
-      // Since portal roots are registered globally, we need this line so that each storybook
-      // story works in isolation.
-      registerPortalRoot(undefined)
       return (
         <ThemeProvider>
           <BaseStyles>
diff --git a/src/stories/Portal.stories.tsx b/src/stories/Portal.stories.tsx
index c65e38d220f..306cec6d8ae 100644
--- a/src/stories/Portal.stories.tsx
+++ b/src/stories/Portal.stories.tsx
@@ -10,9 +10,6 @@ export default {
   component: Portal,
   decorators: [
     Story => {
-      // Since portal roots are registered globally, we need this line so that each storybook
-      // story works in isolation.
-      registerPortalRoot(undefined)
       return (
         <ThemeProvider>
           <BaseStyles>