From 728547d4608e8c601e209ede860493b1986da174 Mon Sep 17 00:00:00 2001
From: jrandolf <101637635+jrandolf@users.noreply.github.com>
Date: Fri, 31 Mar 2023 17:19:59 +0200
Subject: [PATCH] fix: use shadowRoot for tree walker (#9950)

---
 packages/puppeteer-core/src/injected/util.ts |  3 ++-
 test/src/queryhandler.spec.ts                | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/packages/puppeteer-core/src/injected/util.ts b/packages/puppeteer-core/src/injected/util.ts
index 6c4b046d8c5..34fe8f77484 100644
--- a/packages/puppeteer-core/src/injected/util.ts
+++ b/packages/puppeteer-core/src/injected/util.ts
@@ -49,7 +49,8 @@ export function* pierce(root: Node): IterableIterator<Node | ShadowRoot> {
  * @internal
  */
 export function* pierceAll(root: Node): IterableIterator<Node | ShadowRoot> {
-  yield* pierce(root);
+  root = pierce(root).next().value;
+  yield root;
   const walkers = [document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT)];
   for (const walker of walkers) {
     let node: Element | null;
diff --git a/test/src/queryhandler.spec.ts b/test/src/queryhandler.spec.ts
index e5218a30807..ae8c66dece0 100644
--- a/test/src/queryhandler.spec.ts
+++ b/test/src/queryhandler.spec.ts
@@ -404,6 +404,24 @@ describe('Query handler tests', function () {
           })
         ).toBeTruthy();
       }
+      {
+        const elements = await page.$$('#c >>>> div');
+        assert(elements[0], 'Could not find element');
+        expect(
+          await elements[0]?.evaluate(element => {
+            return element.id === 'd';
+          })
+        ).toBeTruthy();
+      }
+      {
+        const elements = await page.$$('#c >>> div');
+        assert(elements[0], 'Could not find element');
+        expect(
+          await elements[0]?.evaluate(element => {
+            return element.id === 'd';
+          })
+        ).toBeTruthy();
+      }
     });
 
     it('should work with text selectors', async () => {