From 2d5572abd8a0f2acec65adfce08dcc2e0f7f70ca Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Sat, 4 Jun 2022 14:07:06 -0700
Subject: [PATCH] fix(ct): return locator to root for fragments (#14639)

fix(fragments): return locator to root for fragments
---
 .../playwright-ct-react/registerSource.mjs    | 24 ++++++++++-----
 .../playwright-ct-svelte/registerSource.mjs   | 23 +++++++++------
 packages/playwright-ct-vue/registerSource.mjs | 23 ++++++---------
 .../playwright-ct-vue2/registerSource.mjs     | 22 +++++---------
 packages/playwright-test/src/mount.ts         | 21 ++++++++++----
 packages/playwright-test/types/component.d.ts | 18 ++++++++----
 .../playwright.ct-react.spec.ts               | 29 ++++++++++++++++++-
 7 files changed, 104 insertions(+), 56 deletions(-)

diff --git a/packages/playwright-ct-react/registerSource.mjs b/packages/playwright-ct-react/registerSource.mjs
index 751de43112..4aa0c99423 100644
--- a/packages/playwright-ct-react/registerSource.mjs
+++ b/packages/playwright-ct-react/registerSource.mjs
@@ -14,18 +14,29 @@
  * limitations under the License.
  */
 
+// @ts-check
 // This file is injected into the registry as text, no dependencies are allowed.
 
 import React from 'react';
 import ReactDOM from 'react-dom';
 
+/** @typedef {import('../playwright-test/types/component').Component} Component */
+/** @typedef {import('react').FunctionComponent} FrameworkComponent */
+
+/** @type {Map<string, FrameworkComponent>} */
 const registry = new Map();
 
+/**
+ * @param {{[key: string]: FrameworkComponent}} components
+ */
 export function register(components) {
   for (const [name, value] of Object.entries(components))
     registry.set(name, value);
 }
 
+/**
+ * @param {Component} component
+ */
 function render(component) {
   let componentFunc = registry.get(component.type);
   if (!componentFunc) {
@@ -41,9 +52,12 @@ function render(component) {
   if (!componentFunc && component.type[0].toUpperCase() === component.type[0])
     throw new Error(`Unregistered component: ${component.type}. Following components are registered: ${[...registry.keys()]}`);
 
-  componentFunc = componentFunc || component.type;
+  const componentFuncOrString = componentFunc || component.type;
 
-  return React.createElement(componentFunc, component.props, ...component.children.map(child => {
+  if (component.kind !== 'jsx')
+    throw new Error('Object mount notation is not supported');
+
+  return React.createElement(componentFuncOrString, component.props, ...component.children.map(child => {
     if (typeof child === 'string')
       return child;
     return render(child);
@@ -55,11 +69,5 @@ function render(component) {
 }
 
 window.playwrightMount = component => {
-  if (!document.getElementById('root')) {
-    const rootElement = document.createElement('div');
-    rootElement.id = 'root';
-    document.body.append(rootElement);
-  }
   ReactDOM.render(render(component), document.getElementById('root'));
-  return '#root > *';
 };
diff --git a/packages/playwright-ct-svelte/registerSource.mjs b/packages/playwright-ct-svelte/registerSource.mjs
index 0248069bac..f3635fbbf4 100644
--- a/packages/playwright-ct-svelte/registerSource.mjs
+++ b/packages/playwright-ct-svelte/registerSource.mjs
@@ -14,21 +14,25 @@
  * limitations under the License.
  */
 
+// @ts-check
+
 // This file is injected into the registry as text, no dependencies are allowed.
 
+/** @typedef {import('../playwright-test/types/component').Component} Component */
+/** @typedef {any} FrameworkComponent */
+
+/** @type {Map<string, FrameworkComponent>} */
 const registry = new Map();
 
+/**
+ * @param {{[key: string]: FrameworkComponent}} components
+ */
 export function register(components) {
   for (const [name, value] of Object.entries(components))
     registry.set(name, value);
 }
 
-window.playwrightMount = component => {
-  if (!document.getElementById('root')) {
-    const rootElement = document.createElement('div');
-    rootElement.id = 'root';
-    document.body.append(rootElement);
-  }
+window.playwrightMount = (component, rootElement) => {
   let componentCtor = registry.get(component.type);
   if (!componentCtor) {
     // Lookup by shorthand.
@@ -43,12 +47,13 @@ window.playwrightMount = component => {
   if (!componentCtor)
     throw new Error(`Unregistered component: ${component.type}. Following components are registered: ${[...registry.keys()]}`);
 
+  if (component.kind !== 'object')
+    throw new Error('JSX mount notation is not supported');
+
   const wrapper = new componentCtor({
-    target: document.getElementById('root'),
+    target: rootElement,
     props: component.options?.props,
   });
-
   for (const [key, listener] of Object.entries(component.options?.on || {}))
     wrapper.$on(key, event => listener(event.detail));
-  return '#root > *';
 };
diff --git a/packages/playwright-ct-vue/registerSource.mjs b/packages/playwright-ct-vue/registerSource.mjs
index 43c2f97835..bee6349cb8 100644
--- a/packages/playwright-ct-vue/registerSource.mjs
+++ b/packages/playwright-ct-vue/registerSource.mjs
@@ -19,13 +19,14 @@
 
 import { createApp, setDevtoolsHook, h } from 'vue';
 
-/** @typedef {import('../playwright-test/types/component').Component} Component */
+/** @typedef {import('@playwright/test/types/component').Component} Component */
+/** @typedef {import('vue').Component} FrameworkComponent */
 
-/** @type { Map<string, import('vue').Component> } */
+/** @type {Map<string, FrameworkComponent>} */
 const registry = new Map();
 
 /**
- * @param {{[key: string]: import('vue').Component}} components
+ * @param {{[key: string]: FrameworkComponent}} components
  */
 export function register(components) {
   for (const [name, value] of Object.entries(components))
@@ -35,7 +36,7 @@ export function register(components) {
 const allListeners = [];
 
 /**
- * @param {Component | string} child 
+ * @param {Component | string} child
  * @returns {import('vue').VNode | string}
  */
 function renderChild(child) {
@@ -43,7 +44,7 @@ function renderChild(child) {
 }
 
 /**
- * @param {Component} component 
+ * @param {Component} component
  * @returns {import('vue').VNode}
  */
 function render(component) {
@@ -70,7 +71,7 @@ function render(component) {
   componentFunc = componentFunc || component.type;
 
   const isVueComponent = componentFunc !== component.type;
- 
+
   /**
    * @type {(import('vue').VNode | string)[]}
    */
@@ -154,16 +155,10 @@ function createDevTools() {
   };
 }
 
-/** @type {any} */ (window).playwrightMount = /** @param {Component} component */ async component => {
-  if (!document.getElementById('root')) {
-    const rootElement = document.createElement('div');
-    rootElement.id = 'root';
-    document.body.append(rootElement);
-  }
+window.playwrightMount = (component, rootElement) => {
   const app = createApp({
     render: () => render(component)
   });
   setDevtoolsHook(createDevTools(), {});
-  app.mount('#root');
-  return '#root > *';
+  app.mount(rootElement);
 };
diff --git a/packages/playwright-ct-vue2/registerSource.mjs b/packages/playwright-ct-vue2/registerSource.mjs
index 70278edae5..83a7313c97 100644
--- a/packages/playwright-ct-vue2/registerSource.mjs
+++ b/packages/playwright-ct-vue2/registerSource.mjs
@@ -21,12 +21,13 @@
 import Vue from 'vue';
 
 /** @typedef {import('../playwright-test/types/component').Component} Component */
+/** @typedef {import('vue').Component} FrameworkComponent */
 
-/** @type { Map<string, import('vue').Component> } */
+/** @type {Map<string, FrameworkComponent>} */
 const registry = new Map();
 
 /**
- * @param {{[key: string]: import('vue').Component}} components
+ * @param {{[key: string]: FrameworkComponent}} components
  */
 export function register(components) {
   for (const [name, value] of Object.entries(components))
@@ -34,8 +35,8 @@ export function register(components) {
 }
 
 /**
- * @param {Component | string} child 
- * @param {import('vue').CreateElement} h 
+ * @param {Component | string} child
+ * @param {import('vue').CreateElement} h
  * @returns {import('vue').VNode | string}
  */
 function renderChild(child, h) {
@@ -43,8 +44,8 @@ function renderChild(child, h) {
 }
 
 /**
- * @param {Component} component 
- * @param {import('vue').CreateElement} h 
+ * @param {Component} component
+ * @param {import('vue').CreateElement} h
  * @returns {import('vue').VNode}
  */
 function render(component, h) {
@@ -133,16 +134,9 @@ function render(component, h) {
   return wrapper;
 }
 
-/** @type {any} */ (window).playwrightMount = /** @param {Component} component */ async component => {
-  let rootElement = document.getElementById('root');
-  if (!rootElement) {
-    rootElement = document.createElement('div');
-    rootElement.id = 'root';
-    document.body.append(rootElement);
-  }
+window.playwrightMount = (component, rootElement) => {
   const mounted = new Vue({
     render: h => render(component, h),
   }).$mount();
   rootElement.appendChild(mounted.$el);
-  return '#root > *';
 };
diff --git a/packages/playwright-test/src/mount.ts b/packages/playwright-test/src/mount.ts
index 372acd2fb6..b76a967b65 100644
--- a/packages/playwright-test/src/mount.ts
+++ b/packages/playwright-test/src/mount.ts
@@ -16,6 +16,7 @@
 
 import { normalizeTraceMode, normalizeVideoMode, shouldCaptureTrace, shouldCaptureVideo } from './index';
 import type { Fixtures, Locator, Page, BrowserContextOptions, PlaywrightTestArgs, PlaywrightTestOptions, PlaywrightWorkerArgs, PlaywrightWorkerOptions, BrowserContext } from './types';
+import type { Component, JsxComponent, ObjectComponentOptions } from '../types/component';
 
 let boundCallbacksForMount: Function[] = [];
 
@@ -65,7 +66,7 @@ export const fixtures: Fixtures<
     },
 
     mount: async ({ page }, use) => {
-      await use(async (component, options) => {
+      await use(async (component: JsxComponent | string, options?: ObjectComponentOptions) => {
         const selector = await (page as any)._wrapApiCall(async () => {
           return await innerMount(page, component, options);
         }, true);
@@ -75,8 +76,8 @@ export const fixtures: Fixtures<
     },
   };
 
-async function innerMount(page: Page, jsxOrType: any, options: any): Promise<string> {
-  let component;
+async function innerMount(page: Page, jsxOrType: JsxComponent | string, options?: ObjectComponentOptions): Promise<string> {
+  let component: Component;
   if (typeof jsxOrType === 'string')
     component = { kind: 'object', type: jsxOrType, options };
   else
@@ -85,7 +86,7 @@ async function innerMount(page: Page, jsxOrType: any, options: any): Promise<str
   wrapFunctions(component, page, boundCallbacksForMount);
 
   // WebKit does not wait for deferred scripts.
-  await page.waitForFunction(() => !!(window as any).playwrightMount);
+  await page.waitForFunction(() => !!window.playwrightMount);
 
   const selector = await page.evaluate(async ({ component }) => {
     const unwrapFunctions = (object: any) => {
@@ -102,7 +103,17 @@ async function innerMount(page: Page, jsxOrType: any, options: any): Promise<str
     };
 
     unwrapFunctions(component);
-    return await (window as any).playwrightMount(component);
+    let rootElement = document.getElementById('root');
+    if (!rootElement) {
+      rootElement = document.createElement('div');
+      rootElement.id = 'root';
+      document.body.appendChild(rootElement);
+    }
+
+    window.playwrightMount(component, rootElement);
+
+    // When mounting fragments, return selector pointing to the root element.
+    return rootElement.childNodes.length > 1 ? '#root' : '#root > *';
   }, { component });
   return selector;
 }
diff --git a/packages/playwright-test/types/component.d.ts b/packages/playwright-test/types/component.d.ts
index 6d628e9f5c..fa0b0812b8 100644
--- a/packages/playwright-test/types/component.d.ts
+++ b/packages/playwright-test/types/component.d.ts
@@ -21,14 +21,22 @@ export type JsxComponent = {
   children: (Component | string)[],
 };
 
+export type ObjectComponentOptions = {
+  props?: { [key: string]: any },
+  slots?: { [key: string]: any },
+  on?: { [key: string]: Function },
+};
+
 export type ObjectComponent = {
   kind: 'object',
   type: string,
-  options?: {
-    props?: { [key: string]: any },
-    slots?: { [key: string]: any },
-    on?: { [key: string]: Function },
-  }
+  options?: ObjectComponentOptions
 };
 
 export type Component = JsxComponent | ObjectComponent;
+
+declare global {
+  interface Window {
+    playwrightMount(component: Component, rootElement: Element): void;
+  }
+}
diff --git a/tests/playwright-test/playwright.ct-react.spec.ts b/tests/playwright-test/playwright.ct-react.spec.ts
index 9252c4e97b..539a609fb6 100644
--- a/tests/playwright-test/playwright.ct-react.spec.ts
+++ b/tests/playwright-test/playwright.ct-react.spec.ts
@@ -262,7 +262,34 @@ test('should work with JSX in variable', async ({ runInlineTest }) => {
 
       test('pass button', async ({ mount }) => {
         const component = await mount(button);
-        await expect(component).toHaveText('Button'); 
+        await expect(component).toHaveText('Button');
+      });
+    `,
+  }, { workers: 1 });
+
+  expect(result.exitCode).toBe(0);
+  expect(result.passed).toBe(1);
+});
+
+test('should return root locator for fragments', async ({ runInlineTest }) => {
+  const result = await runInlineTest({
+    'playwright/index.html': `<script type="module" src="/playwright/index.js"></script>`,
+    'playwright/index.js': `//@no-header`,
+
+    'src/button.jsx': `
+      //@no-header
+      export const Button = () => <><h1>Header</h1><button>Button</button></>;
+    `,
+
+    'src/button.test.jsx': `
+      //@no-header
+      import { test, expect } from '@playwright/experimental-ct-react';
+      import { Button } from './button';
+
+      test('pass button', async ({ mount }) => {
+        const component = await mount(<Button></Button>);
+        await expect(component).toContainText('Header');
+        await expect(component).toContainText('Button');
       });
     `,
   }, { workers: 1 });