From 1edc1d8c4b64c396bb8343528c1a8474054a515b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=8C=E8=B4=A7=E6=9C=BA=E5=99=A8=E4=BA=BA?=
 <smith3816@gmail.com>
Date: Tue, 27 Sep 2022 12:08:22 +0800
Subject: [PATCH] feat: Enable disable port in test env

---
 src/Portal.tsx                            | 30 ++++++++++++++++++-----
 src/index.tsx                             |  2 ++
 src/mock.ts                               |  8 ++++++
 tests/__snapshots__/testEnv.test.tsx.snap | 13 ++++++++++
 tests/setupAfterEnv.ts                    |  1 +
 tests/ssr.test.tsx                        | 19 ++++++++++++++
 tests/testEnv.test.tsx                    | 21 ++++++++++++++++
 7 files changed, 88 insertions(+), 6 deletions(-)
 create mode 100644 src/mock.ts
 create mode 100644 tests/__snapshots__/testEnv.test.tsx.snap
 create mode 100644 tests/setupAfterEnv.ts
 create mode 100644 tests/ssr.test.tsx
 create mode 100644 tests/testEnv.test.tsx

diff --git a/src/Portal.tsx b/src/Portal.tsx
index 372046b..8699590 100644
--- a/src/Portal.tsx
+++ b/src/Portal.tsx
@@ -1,15 +1,18 @@
 import * as React from 'react';
 import { createPortal } from 'react-dom';
+import canUseDom from 'rc-util/lib/Dom/canUseDom';
 import OrderContext from './Context';
 import useDom from './useDom';
 import useScrollLocker from './useScrollLocker';
+import { inlineMock } from './mock';
 
-// ZombieJ: Since React 18 strict mode logic change.
-//          We should rewrite for compatible.
+export type ContainerType = Element | DocumentFragment;
+
+export type GetContainer = string | ContainerType | (() => ContainerType);
 
 export interface PortalProps {
   /** Customize container element. Default will create a div in document.body when `open` */
-  getContainer?: () => Element | DocumentFragment;
+  getContainer?: GetContainer;
   children?: React.ReactNode;
   /** Show the portal children */
   open?: boolean;
@@ -17,6 +20,20 @@ export interface PortalProps {
   autoLock?: boolean;
 }
 
+const getPortalContainer = (getContainer: GetContainer) => {
+  if (!canUseDom()) {
+    return null;
+  }
+
+  if (typeof getContainer === 'string') {
+    return document.querySelector(getContainer);
+  }
+  if (typeof getContainer === 'function') {
+    return getContainer();
+  }
+  return getContainer;
+};
+
 export default function Portal(props: PortalProps) {
   const { open, autoLock, getContainer, children } = props;
 
@@ -30,7 +47,7 @@ export default function Portal(props: PortalProps) {
   }, [open]);
 
   // ======================== Container ========================
-  const customizeContainer = getContainer?.();
+  const customizeContainer = getPortalContainer(getContainer);
 
   const [defaultContainer, queueCreate] = useDom(
     mergedRender && !customizeContainer,
@@ -39,13 +56,14 @@ export default function Portal(props: PortalProps) {
 
   // ========================= Render ==========================
   // Do not render when nothing need render
-  if (!mergedRender) {
+  if (!mergedRender || !canUseDom()) {
     return null;
   }
 
+  console.log(inlineMock());
   return (
     <OrderContext.Provider value={queueCreate}>
-      {createPortal(children, mergedContainer)}
+      {inlineMock() ? children : createPortal(children, mergedContainer)}
     </OrderContext.Provider>
   );
 }
diff --git a/src/index.tsx b/src/index.tsx
index 3cb0f84..34fa676 100644
--- a/src/index.tsx
+++ b/src/index.tsx
@@ -1,6 +1,8 @@
 import Portal from './Portal';
 import type { PortalProps } from './Portal';
+import { inlineMock } from './mock';
 
 export type { PortalProps };
+export { inlineMock };
 
 export default Portal;
diff --git a/src/mock.ts b/src/mock.ts
new file mode 100644
index 0000000..6fefe84
--- /dev/null
+++ b/src/mock.ts
@@ -0,0 +1,8 @@
+export let inline = false;
+
+export function inlineMock(nextInline?: boolean) {
+  if (typeof nextInline === 'boolean') {
+    inline = nextInline;
+  }
+  return inline;
+}
diff --git a/tests/__snapshots__/testEnv.test.tsx.snap b/tests/__snapshots__/testEnv.test.tsx.snap
new file mode 100644
index 0000000..91bdcb4
--- /dev/null
+++ b/tests/__snapshots__/testEnv.test.tsx.snap
@@ -0,0 +1,13 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`Test Env inlineMock 1`] = `
+<div>
+  Start
+  <div
+    class="bamboo"
+  >
+    Hello World
+  </div>
+  End
+</div>
+`;
diff --git a/tests/setupAfterEnv.ts b/tests/setupAfterEnv.ts
new file mode 100644
index 0000000..7b0828b
--- /dev/null
+++ b/tests/setupAfterEnv.ts
@@ -0,0 +1 @@
+import '@testing-library/jest-dom';
diff --git a/tests/ssr.test.tsx b/tests/ssr.test.tsx
new file mode 100644
index 0000000..20f3d6c
--- /dev/null
+++ b/tests/ssr.test.tsx
@@ -0,0 +1,19 @@
+import React from 'react';
+import { render } from '@testing-library/react';
+import Portal from '../src';
+
+jest.mock('rc-util/lib/Dom/canUseDom', () => () => false);
+
+describe('SSR', () => {
+  it('No Render in SSR', () => {
+    const { unmount } = render(
+      <Portal open>
+        <div className="bamboo">Hello World</div>
+      </Portal>,
+    );
+
+    expect(document.querySelector('.bamboo')).toBeFalsy();
+
+    unmount();
+  });
+});
diff --git a/tests/testEnv.test.tsx b/tests/testEnv.test.tsx
new file mode 100644
index 0000000..4cf04e6
--- /dev/null
+++ b/tests/testEnv.test.tsx
@@ -0,0 +1,21 @@
+import React from 'react';
+import { render } from '@testing-library/react';
+import Portal, { inlineMock } from '../src';
+
+describe('Test Env', () => {
+  it('inlineMock', () => {
+    inlineMock(true);
+
+    const { container } = render(
+      <>
+        Start
+        <Portal open>
+          <div className="bamboo">Hello World</div>
+        </Portal>
+        End
+      </>,
+    );
+
+    expect(container).toMatchSnapshot();
+  });
+});