Additional test and cleaning up stale values

Author: mattrunyon

packages/components/src/spectrum/TabPanels.test.tsx

@@ -1,5 +1,5 @@
-import React, { useState } from 'react';
-import { render, screen } from '@testing-library/react';
+import React, { useEffect, useState } from 'react';
+import { render, screen, waitFor } from '@testing-library/react';
 import {
   defaultTheme,
   Item,
@@ -20,6 +20,22 @@ function Counter({ label }: { label: string }) {
   );
 }
 
+function OnMountUnmount({
+  onMount,
+  onUnmount,
+}: {
+  onMount: () => void;
+  onUnmount: () => void;
+}) {
+  useEffect(() => {
+    onMount();
+    return () => {
+      onUnmount();
+    };
+  }, [onMount, onUnmount]);
+  return null;
+}
+
 describe('TabPanels', () => {
   it('should not persist panel state by default when switching tabs', () => {
     render(
@@ -152,4 +168,51 @@ describe('TabPanels', () => {
       'background-color: red'
     );
   });
+
+  it('should still unmount a panel that is not in the tree when using keepMounted', () => {
+    const onMount = jest.fn();
+    const onUnmount = jest.fn();
+    const { rerender } = render(
+      <Provider theme={defaultTheme}>
+        <Tabs aria-label="test">
+          <TabList>
+            <Item key="1">Tab 1</Item>
+            <Item key="2">Tab 2</Item>
+          </TabList>
+          <DHCTabPanels>
+            <Item key="1">
+              <Counter label="foo" />
+            </Item>
+            <Item key="2">
+              <OnMountUnmount onMount={onMount} onUnmount={onUnmount} />
+            </Item>
+          </DHCTabPanels>
+        </Tabs>
+      </Provider>
+    );
+
+    waitFor(() => expect(onMount).toHaveBeenCalledTimes(1));
+    expect(onUnmount).toHaveBeenCalledTimes(0);
+
+    rerender(
+      <Provider theme={defaultTheme}>
+        <Tabs aria-label="test">
+          <TabList>
+            <Item key="1">Tab 1</Item>
+            <Item key="2">Tab 2</Item>
+          </TabList>
+          <DHCTabPanels>
+            <Item key="1">
+              <Counter label="foo" />
+            </Item>
+            <Item key="2">
+              <Counter label="bar" />
+            </Item>
+          </DHCTabPanels>
+        </Tabs>
+      </Provider>
+    );
+
+    waitFor(() => expect(onUnmount).toHaveBeenCalledTimes(1));
+  });
 });

packages/components/src/spectrum/TabPanels.tsx

@@ -36,7 +36,9 @@ export function DHCTabPanels<T extends object>(
 
   const portalNodes = useMemo(() => {
     const nodes: JSX.Element[] = [];
+    const nextNodeMap = new Map<Key, HtmlPortalNode>(); // Keep track of the portals we use so we can clean up stale portals
     if (!keepMounted) {
+      portalNodeMap.current = nextNodeMap;
       return nodes;
     }
     React.Children.forEach(children, child => {
@@ -53,15 +55,17 @@ export function DHCTabPanels<T extends object>(
             style: 'display: contents',
           },
         });
-        portalNodeMap.current.set(child.key, portal);
       }
+      nextNodeMap.set(child.key, portal);
       nodes.push(
         <InPortal node={portal} key={child.key}>
           {child.props.children}
         </InPortal>
       );
     });
 
+    portalNodeMap.current = nextNodeMap;
+
     return nodes;
   }, [children, keepMounted]);