feat: support multiple roots in React and Vue selectors (#8313)

Fixes #8230

src/server/injected/reactSelectorEngine.ts

@@ -132,30 +132,36 @@ function filterComponentsTree(treeNode: ComponentNode, searchFn: (node: Componen
   return result;
 }
 
-function findReactRoot(): ReactVNode | undefined {
+function findReactRoots(): ReactVNode[] {
-  const walker = document.createTreeWalker(document);
+  const roots: ReactVNode[] = [];
-  while (walker.nextNode()) {
+  const walker = document.createTreeWalker(document, NodeFilter.SHOW_ELEMENT,   {
+    acceptNode: function(node) {
       // @see https://github.com/baruchvlz/resq/blob/5c15a5e04d3f7174087248f5a158c3d6dcc1ec72/src/utils.js#L329
-    if (walker.currentNode.hasOwnProperty('_reactRootContainer'))
+      if (node.hasOwnProperty('_reactRootContainer')) {
-      return (walker.currentNode as any)._reactRootContainer._internalRoot.current;
+        roots.push((node as any)._reactRootContainer._internalRoot.current);
-    for (const key of Object.keys(walker.currentNode)) {
+        return NodeFilter.FILTER_REJECT;
+      }
+      for (const key of Object.keys(node)) {
         // @see https://github.com/baruchvlz/resq/blob/5c15a5e04d3f7174087248f5a158c3d6dcc1ec72/src/utils.js#L334
-      if (key.startsWith('__reactInternalInstance') || key.startsWith('__reactFiber'))
+        if (key.startsWith('__reactInternalInstance') || key.startsWith('__reactFiber')) {
-        return (walker.currentNode as any)[key];
+          roots.push((node as any)[key]);
+          return NodeFilter.FILTER_REJECT;
         }
       }
-  return undefined;
+      return NodeFilter.FILTER_ACCEPT;
+    }
+  });
+  while (walker.nextNode());
+  return roots;
 }
 
 export const ReactEngine: SelectorEngine = {
   queryAll(scope: SelectorRoot, selector: string): Element[] {
     const {name, attributes} = parseComponentSelector(selector);
 
-    const reactRoot = findReactRoot();
+    const reactRoots = findReactRoots();
-    if (!reactRoot)
+    const trees = reactRoots.map(reactRoot => buildComponentsTree(reactRoot));
-      return [];
+    const treeNodes = trees.map(tree => filterComponentsTree(tree, treeNode => {
-    const tree = buildComponentsTree(reactRoot);
-    const treeNodes = filterComponentsTree(tree, treeNode => {
       if (name && treeNode.name !== name)
         return false;
       if (treeNode.rootElements.some(domNode => !scope.contains(domNode)))
@@ -165,7 +171,7 @@ export const ReactEngine: SelectorEngine = {
           return false;
       }
       return true;
-    });
+    })).flat();
     const allRootElements: Set<Element> = new Set();
     for (const treeNode of treeNodes) {
       for (const domNode of treeNode.rootElements)

src/server/injected/vueSelectorEngine.ts

@@ -204,27 +204,34 @@ function filterComponentsTree(treeNode: ComponentNode, searchFn: (node: Componen
   return result;
 }
 
-function findVueRoot(): undefined|{version: number, root: VueVNode} {
+type VueRoot = {version: number, root: VueVNode};
-  const walker = document.createTreeWalker(document);
+function findVueRoots(): VueRoot[] {
-  while (walker.nextNode()) {
+  const roots: VueRoot[] = [];
+  const walker = document.createTreeWalker(document, NodeFilter.SHOW_ELEMENT, {
+    acceptNode: function(node) {
       // Vue3 root
-    if ((walker.currentNode as any)._vnode && (walker.currentNode as any)._vnode.component)
+      if ((node as any)._vnode && (node as any)._vnode.component) {
-      return {root: (walker.currentNode as any)._vnode.component, version: 3};
+        roots.push({root: (node as any)._vnode.component, version: 3});
-    // Vue2 root
+        return NodeFilter.FILTER_REJECT;
-    if ((walker.currentNode as any).__vue__)
-      return {root: (walker.currentNode as any).__vue__, version: 2};
       }
-  return undefined;
+      // Vue2 root
+      if ((node as any).__vue__) {
+        roots.push({root: (node as any).__vue__, version: 2});
+        return NodeFilter.FILTER_REJECT;
+      }
+      return NodeFilter.FILTER_ACCEPT;
+    }
+  });
+  while (walker.nextNode());
+  return roots;
 }
 
 export const VueEngine: SelectorEngine = {
   queryAll(scope: SelectorRoot, selector: string): Element[] {
     const {name, attributes} = parseComponentSelector(selector);
-    const vueRoot = findVueRoot();
+    const vueRoots = findVueRoots();
-    if (!vueRoot)
+    const trees = vueRoots.map(vueRoot => vueRoot.version === 3 ? buildComponentsTreeVue3(vueRoot.root) : buildComponentsTreeVue2(vueRoot.root));
-      return [];
+    const treeNodes = trees.map(tree => filterComponentsTree(tree, treeNode => {
-    const tree = vueRoot.version === 3 ? buildComponentsTreeVue3(vueRoot.root) : buildComponentsTreeVue2(vueRoot.root);
-    const treeNodes = filterComponentsTree(tree, treeNode => {
       if (name && treeNode.name !== name)
         return false;
       if (treeNode.rootElements.some(rootElement => !scope.contains(rootElement)))
@@ -233,9 +240,8 @@ export const VueEngine: SelectorEngine = {
         if (!checkComponentAttribute(treeNode.props, attr))
           return false;
       }
-
       return true;
-    });
+    })).flat();
     const allRootElements: Set<Element> = new Set();
     for (const treeNode of treeNodes) {
       for (const rootElement of treeNode.rootElements)

tests/assets/reading-list/react15.html

@@ -101,9 +101,7 @@ class App extends React.Component {
   }
 }
 
-ReactDOM.render(
+window.mountApp = element => ReactDOM.render(e(App, null, null), element);
-  e(App, null, null),
+window.mountApp(document.getElementById('root'));
-  document.getElementById('root'),
-);
 
 </script>

tests/assets/reading-list/react16.html

@@ -99,9 +99,7 @@ class App extends React.Component {
   }
 }
 
-ReactDOM.render(
+window.mountApp = element => ReactDOM.render(e(App, null, null), element);
-  e(App, null, null),
+window.mountApp(document.getElementById('root'));
-  document.getElementById('root'),
-);
 
 </script>

tests/assets/reading-list/react17.html

@@ -93,9 +93,7 @@ class App extends React.Component {
   }
 }
 
-ReactDOM.render(
+window.mountApp = element => ReactDOM.render(e(App, null, null), element);
-  e(App, null, null),
+window.mountApp(document.getElementById('root'));
-  document.getElementById('root'),
-);
 
 </script>

tests/assets/reading-list/vue2.html

@@ -1,7 +1,7 @@
 <link rel=stylesheet href='./style.css'>
 <script src="./vue_2.6.14.js"></script>
 
-<div id=root></div>
+<div id=root><div></div></div>
 
 <script>
 
@@ -82,9 +82,8 @@ Vue.component('button-grid', {
   }
 });
 
-
+window.mountApp = element => new Vue({
-new Vue({
+  el: element,
-  el: '#root',
 
   template: `
     <div>
@@ -112,5 +111,6 @@ new Vue({
     }
   },
 });
+window.mountApp(document.querySelector('#root div'));
 
 </script>

tests/assets/reading-list/vue3.html

@@ -1,10 +1,11 @@
 <link rel=stylesheet href='./style.css'>
 <script src="./vue_3.1.5.js"></script>
 
-<div id=root></div>
+<div id=root><div></div></div>
 
 <script>
 
+window.mountApp = element => {
   const app = Vue.createApp({
     template: `
       <app-header :bookCount='books.length'></app-header>
@@ -106,7 +107,9 @@ app.component('button-grid', {
       return buttons;
     }
   });
+  app.mount(element);
+}
 
-app.mount('#root');
+window.mountApp(document.querySelector('#root div'));
 
 </script>

tests/page/selectors-react.spec.ts

@@ -104,6 +104,27 @@ for (const [name, url] of Object.entries(reacts)) {
     it('should support truthy querying', async ({page}) => {
       expect(await page.$$eval(`_react=ColorButton[enabled]`, els => els.length)).toBe(5);
     });
+
+    it('should work with multiroot react', async ({page}) => {
+      await it.step('mount second root', async () => {
+        await expect(page.locator(`_react=BookItem`)).toHaveCount(3);
+        await page.evaluate(() => {
+          const anotherRoot = document.createElement('div');
+          anotherRoot.id = 'root2';
+          document.body.append(anotherRoot);
+          // @ts-ignore
+          window.mountApp(anotherRoot);
+        });
+        await expect(page.locator(`_react=BookItem`)).toHaveCount(6);
+      });
+
+      await it.step('add a new book to second root', async () => {
+        await page.locator('#root2 input').fill('newbook');
+        await page.locator('#root2 >> text=new book').click();
+        await expect(page.locator('css=#root >> _react=BookItem')).toHaveCount(3);
+        await expect(page.locator('css=#root2 >> _react=BookItem')).toHaveCount(4);
+      });
+    });
   });
 }
 

tests/page/selectors-vue.spec.ts

@@ -101,6 +101,28 @@ for (const [name, url] of Object.entries(vues)) {
     it('should support truthy querying', async ({page}) => {
       expect(await page.$$eval(`_vue=color-button[enabled]`, els => els.length)).toBe(5);
     });
+
+    it('should work with multiroot react', async ({page}) => {
+      await it.step('mount second root', async () => {
+        await expect(page.locator(`_vue=book-item`)).toHaveCount(3);
+        await page.evaluate(() => {
+          const anotherRoot = document.createElement('div');
+          anotherRoot.id = 'root2';
+          anotherRoot.append(document.createElement('div'));
+          document.body.append(anotherRoot);
+          // @ts-ignore
+          window.mountApp(anotherRoot.querySelector('div'));
+        });
+        await expect(page.locator(`_vue=book-item`)).toHaveCount(6);
+      });
+
+      await it.step('add a new book to second root', async () => {
+        await page.locator('#root2 input').fill('newbook');
+        await page.locator('#root2 >> text=new book').click();
+        await expect(page.locator('css=#root >> _vue=book-item')).toHaveCount(3);
+        await expect(page.locator('css=#root2 >> _vue=book-item')).toHaveCount(4);
+      });
+    });
   });
 }