From 7e6e5f070630cd63c723b7390b329498abc2f9c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Neves?= <joaomiguelneves@gmail.com>
Date: Wed, 31 May 2023 01:14:47 +0100
Subject: [PATCH] feat: Support React forwards refs and memo (#23262)

This PR fixes the react selector behavior to support components that are
wrapped by the memo or forwardRef React builtin functions.

Previously these components couldn't be selected. This PR fixes that
behavior, enabling selecting those components.

Current behavior:
```
const Foo = memo(() => <div id="foo_component" />);
Foo.displayName = "Foo";
...
playwright.$("_react=Foo") -> undefined
```

Fixed behavior:
```
const Foo = memo(() => <div id="foo_component" />);
Foo.displayName = "Foo";
...
playwright.$("_react=Foo") -> <div id ="foo_component" />
```
---
 .../src/server/injected/reactSelectorEngine.ts | 18 ++++++++++++++----
 tests/assets/reading-list/react17.html         |  6 ++++--
 tests/assets/reading-list/react18.html         |  6 ++++--
 tests/page/selectors-react.spec.ts             |  5 +++++
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/packages/playwright-core/src/server/injected/reactSelectorEngine.ts b/packages/playwright-core/src/server/injected/reactSelectorEngine.ts
index 1fad96b0c5..2d0e51f0b2 100644
--- a/packages/playwright-core/src/server/injected/reactSelectorEngine.ts
+++ b/packages/playwright-core/src/server/injected/reactSelectorEngine.ts
@@ -43,13 +43,23 @@ type ReactVNode = {
   _renderedChildren?: any[],
 };
 
+function getFunctionComponentName(component: any) {
+  return component.displayName || component.name || 'Anonymous';
+}
+
 function getComponentName(reactElement: ReactVNode): string {
   // React 16+
   // @see https://github.com/baruchvlz/resq/blob/5c15a5e04d3f7174087248f5a158c3d6dcc1ec72/src/utils.js#L16
-  if (typeof reactElement.type === 'function')
-    return reactElement.type.displayName || reactElement.type.name || 'Anonymous';
-  if (typeof reactElement.type === 'string')
-    return reactElement.type;
+  if (reactElement.type) {
+    switch (typeof reactElement.type) {
+      case 'function':
+        return getFunctionComponentName(reactElement.type);
+      case 'string':
+        return reactElement.type;
+      case 'object': // support memo and forwardRef
+        return reactElement.type.displayName || (reactElement.type.render ? getFunctionComponentName(reactElement.type.render) : '');
+    }
+  }
 
   // React 15
   // @see https://github.com/facebook/react/blob/2edf449803378b5c58168727d4f123de3ba5d37f/packages/react-devtools-shared/src/backend/legacy/renderer.js#L59
diff --git a/tests/assets/reading-list/react17.html b/tests/assets/reading-list/react17.html
index e9a8d7eeaf..9cb8eaba2f 100644
--- a/tests/assets/reading-list/react17.html
+++ b/tests/assets/reading-list/react17.html
@@ -38,7 +38,7 @@ function ColorButton (props) {
   return e('button', {className: props.color, disabled: !props.enabled}, 'button ' + props.nested.index);
 }
 
-function ButtonGrid() {
+const ButtonGrid = React.memo(function() {
   const buttons = [];
   for (let i = 0; i < 9; ++i) {
     buttons.push(e(ColorButton, {
@@ -51,7 +51,9 @@ function ButtonGrid() {
     }, null));
   };
   return e(React.Fragment, null, ...buttons);
-}
+});
+
+ButtonGrid.displayName = "ButtonGrid";
 
 class BookItem extends React.Component {
   render() {
diff --git a/tests/assets/reading-list/react18.html b/tests/assets/reading-list/react18.html
index 30e0d26e73..a521c620d0 100644
--- a/tests/assets/reading-list/react18.html
+++ b/tests/assets/reading-list/react18.html
@@ -38,7 +38,7 @@ function ColorButton (props) {
   return e('button', {className: props.color, disabled: !props.enabled}, 'button ' + props.nested.index);
 }
 
-function ButtonGrid() {
+const ButtonGrid = React.memo(function() {
   const buttons = [];
   for (let i = 0; i < 9; ++i) {
     buttons.push(e(ColorButton, {
@@ -51,7 +51,9 @@ function ButtonGrid() {
     }, null));
   };
   return e(React.Fragment, null, ...buttons);
-}
+});
+
+ButtonGrid.displayName = "ButtonGrid";
 
 class BookItem extends React.Component {
   render() {
diff --git a/tests/page/selectors-react.spec.ts b/tests/page/selectors-react.spec.ts
index 9ad71532aa..ce47f1bb25 100644
--- a/tests/page/selectors-react.spec.ts
+++ b/tests/page/selectors-react.spec.ts
@@ -124,6 +124,11 @@ for (const [name, url] of Object.entries(reacts)) {
       await expect(page.locator(`_react=BookItem`)).toHaveCount(6);
     });
 
+    it('should work with react memo', async ({ page }) => {
+      it.skip(name === 'react15' || name === 'react16', 'Class components dont support memo');
+      await expect(page.locator(`_react=ButtonGrid`)).toHaveCount(9);
+    });
+
     it('should work with multiroot react', async ({ page }) => {
       await it.step('mount second root', async () => {
         await expect(page.locator(`_react=BookItem`)).toHaveCount(3);