From d1dd68a9259c48529ea85919f4ec412eec1d5b3c Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Thu, 18 Jan 2024 13:12:16 -0800
Subject: [PATCH] fix(ct-react): do not reset mount hooks upon update

---
 .../playwright-ct-react/registerSource.mjs    | 43 +++++++++++--------
 .../playwright-ct-react17/registerSource.mjs  | 30 ++++++++++---
 tests/components/ct-react-vite/src/App.tsx    |  3 +-
 .../ct-react-vite/tests/react-router.spec.tsx | 12 ++++++
 tests/components/ct-react17/src/App.tsx       |  3 +-
 .../ct-react17/tests/react-router.spec.tsx    | 12 ++++++
 6 files changed, 78 insertions(+), 25 deletions(-)

diff --git a/packages/playwright-ct-react/registerSource.mjs b/packages/playwright-ct-react/registerSource.mjs
index 0a30836824035..d5e617785a805 100644
--- a/packages/playwright-ct-react/registerSource.mjs
+++ b/packages/playwright-ct-react/registerSource.mjs
@@ -21,7 +21,7 @@ import __pwReact from 'react';
 import { createRoot as __pwCreateRoot } from 'react-dom/client';
 /** @typedef {import('../playwright-ct-core/types/component').JsxComponent} JsxComponent */
 
-/** @type {Map<Element, import('react-dom/client').Root>} */
+/** @type {Map<Element, { root: import('react-dom/client').Root, setRenderer: (renderer: any) => void }>} */
 const __pwRootRegistry = new Map();
 
 /**
@@ -48,33 +48,41 @@ function __pwRender(value) {
 window.playwrightMount = async (component, rootElement, hooksConfig) => {
   if (!isJsxComponent(component))
     throw new Error('Object mount notation is not supported');
-
-  let App = () => __pwRender(component);
-  for (const hook of window.__pw_hooks_before_mount || []) {
-    const wrapper = await hook({ App, hooksConfig });
-    if (wrapper)
-      App = () => wrapper;
-  }
-
   if (__pwRootRegistry.has(rootElement)) {
     throw new Error(
         'Attempting to mount a component into an container that already has a React root'
     );
   }
+
   const root = __pwCreateRoot(rootElement);
-  __pwRootRegistry.set(rootElement, root);
-  root.render(App());
+  const entry = { root, setRenderer: () => undefined };
+  __pwRootRegistry.set(rootElement, entry);
+
+  const App = () => {
+    /** @type {any} */
+    const [renderer, setRenderer] = __pwReact.useState(() => __pwRender(component));
+    entry.setRenderer = setRenderer;
+    return renderer;
+  };
+  let AppWrapper = App;
+  for (const hook of window.__pw_hooks_before_mount || []) {
+    const wrapper = await hook({ App: AppWrapper, hooksConfig });
+    if (wrapper)
+      AppWrapper = () => wrapper;
+  }
+
+  root.render(__pwReact.createElement(AppWrapper));
 
   for (const hook of window.__pw_hooks_after_mount || [])
     await hook({ hooksConfig });
 };
 
 window.playwrightUnmount = async rootElement => {
-  const root = __pwRootRegistry.get(rootElement);
-  if (root === undefined)
+  const entry = __pwRootRegistry.get(rootElement);
+  if (!entry)
     throw new Error('Component was not mounted');
 
-  root.unmount();
+  entry.root.unmount();
   __pwRootRegistry.delete(rootElement);
 };
 
@@ -82,9 +90,8 @@ window.playwrightUpdate = async (rootElement, component) => {
   if (!isJsxComponent(component))
     throw new Error('Object mount notation is not supported');
 
-  const root = __pwRootRegistry.get(rootElement);
-  if (root === undefined)
+  const entry = __pwRootRegistry.get(rootElement);
+  if (!entry)
     throw new Error('Component was not mounted');
-
-  root.render(__pwRender(component));
+  entry.setRenderer(() => __pwRender(component));
 };
diff --git a/packages/playwright-ct-react17/registerSource.mjs b/packages/playwright-ct-react17/registerSource.mjs
index 9363b4bb22d52..2a13152dceedf 100644
--- a/packages/playwright-ct-react17/registerSource.mjs
+++ b/packages/playwright-ct-react17/registerSource.mjs
@@ -21,6 +21,9 @@ import __pwReact from 'react';
 import __pwReactDOM from 'react-dom';
 /** @typedef {import('../playwright-ct-core/types/component').JsxComponent} JsxComponent */
 
+/** @type {Map<Element, { setRenderer: (renderer: any) => void }>} */
+const __pwRootRegistry = new Map();
+
 /**
  * @param {any} component
  * @returns {component is JsxComponent}
@@ -45,15 +48,29 @@ function __pwRender(value) {
 window.playwrightMount = async (component, rootElement, hooksConfig) => {
   if (!isJsxComponent(component))
     throw new Error('Object mount notation is not supported');
+  if (__pwRootRegistry.has(rootElement)) {
+    throw new Error(
+        'Attempting to mount a component into an container that already has a React root'
+    );
+  }
 
-  let App = () => __pwRender(component);
+  const entry = { setRenderer: () => undefined };
+  __pwRootRegistry.set(rootElement, entry);
+
+  const App = () => {
+    /** @type {any} */
+    const [renderer, setRenderer] = __pwReact.useState(() => __pwRender(component));
+    entry.setRenderer = setRenderer;
+    return renderer;
+  };
+  let AppWrapper = App;
   for (const hook of window.__pw_hooks_before_mount || []) {
-    const wrapper = await hook({ App, hooksConfig });
+    const wrapper = await hook({ App: AppWrapper, hooksConfig });
     if (wrapper)
-      App = () => wrapper;
+      AppWrapper = () => wrapper;
   }
 
-  __pwReactDOM.render(App(), rootElement);
+  __pwReactDOM.render(__pwReact.createElement(AppWrapper), rootElement);
 
   for (const hook of window.__pw_hooks_after_mount || [])
     await hook({ hooksConfig });
@@ -68,5 +85,8 @@ window.playwrightUpdate = async (rootElement, component) => {
   if (!isJsxComponent(component))
     throw new Error('Object mount notation is not supported');
 
-  __pwReactDOM.render(__pwRender(component), rootElement);
+  const entry = __pwRootRegistry.get(rootElement);
+  if (!entry)
+    throw new Error('Component was not mounted');
+  entry.setRenderer(() => __pwRender(component));
 };
diff --git a/tests/components/ct-react-vite/src/App.tsx b/tests/components/ct-react-vite/src/App.tsx
index 5f4afe2236dca..fd004390f3b1f 100644
--- a/tests/components/ct-react-vite/src/App.tsx
+++ b/tests/components/ct-react-vite/src/App.tsx
@@ -3,10 +3,11 @@ import logo from './assets/logo.svg';
 import LoginPage from './pages/LoginPage';
 import DashboardPage from './pages/DashboardPage';
 
-export default function App() {
+export default function App({ title }: { title?: string }) {
   return <>
     <header>
       <img src={logo} alt="logo" width={125} height={125} />
+      {title && <h1>{title}</h1>}
       <Link to="/">Login</Link>
       <Link to="/dashboard">Dashboard</Link>
     </header>
diff --git a/tests/components/ct-react-vite/tests/react-router.spec.tsx b/tests/components/ct-react-vite/tests/react-router.spec.tsx
index 4d29a9dc8ade7..939154e9d49e3 100644
--- a/tests/components/ct-react-vite/tests/react-router.spec.tsx
+++ b/tests/components/ct-react-vite/tests/react-router.spec.tsx
@@ -12,3 +12,15 @@ test('navigate to a page by clicking a link', async ({ page, mount }) => {
   await expect(component.getByRole('main')).toHaveText('Dashboard');
   await expect(page).toHaveURL('/dashboard');
 });
+
+test('update should not reset mount hooks', async ({ page, mount }) => {
+  const component = await mount<HooksConfig>(<App title='before'/>, {
+    hooksConfig: { routing: true },
+  });
+  await expect(component.getByRole('heading')).toHaveText('before');
+  await expect(component.getByRole('main')).toHaveText('Login');
+
+  await component.update(<App title='after'/>);
+  await expect(component.getByRole('heading')).toHaveText('after');
+  await expect(component.getByRole('main')).toHaveText('Login');
+});
diff --git a/tests/components/ct-react17/src/App.tsx b/tests/components/ct-react17/src/App.tsx
index 5f4afe2236dca..fd004390f3b1f 100644
--- a/tests/components/ct-react17/src/App.tsx
+++ b/tests/components/ct-react17/src/App.tsx
@@ -3,10 +3,11 @@ import logo from './assets/logo.svg';
 import LoginPage from './pages/LoginPage';
 import DashboardPage from './pages/DashboardPage';
 
-export default function App() {
+export default function App({ title }: { title?: string }) {
   return <>
     <header>
       <img src={logo} alt="logo" width={125} height={125} />
+      {title && <h1>{title}</h1>}
       <Link to="/">Login</Link>
       <Link to="/dashboard">Dashboard</Link>
     </header>
diff --git a/tests/components/ct-react17/tests/react-router.spec.tsx b/tests/components/ct-react17/tests/react-router.spec.tsx
index 98004cef1d5a6..73c4a50c7e4f5 100644
--- a/tests/components/ct-react17/tests/react-router.spec.tsx
+++ b/tests/components/ct-react17/tests/react-router.spec.tsx
@@ -12,3 +12,15 @@ test('navigate to a page by clicking a link', async ({ page, mount }) => {
   await expect(component.getByRole('main')).toHaveText('Dashboard');
   await expect(page).toHaveURL('/dashboard');
 });
+
+test('update should not reset mount hooks', async ({ page, mount }) => {
+  const component = await mount<HooksConfig>(<App title='before'/>, {
+    hooksConfig: { routing: true },
+  });
+  await expect(component.getByRole('heading')).toHaveText('before');
+  await expect(component.getByRole('main')).toHaveText('Login');
+
+  await component.update(<App title='after'/>);
+  await expect(component.getByRole('heading')).toHaveText('after');
+  await expect(component.getByRole('main')).toHaveText('Login');
+});