From 8fabe328006b9741eba53fa4f7f319d3b3286fe2 Mon Sep 17 00:00:00 2001
From: Johan Bay <jobay@google.com>
Date: Tue, 13 Oct 2020 11:05:47 +0200
Subject: [PATCH] feat(queryhandler): add built-in pierce handler (#6509)

Adds a handler 'pierce' that pierces shadow roots while querying.
---
 src/common/QueryHandler.ts | 55 +++++++++++++++++++++++++++++++++++++-
 test/queryselector.spec.ts | 39 +++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/src/common/QueryHandler.ts b/src/common/QueryHandler.ts
index 71386683..70c0d908 100644
--- a/src/common/QueryHandler.ts
+++ b/src/common/QueryHandler.ts
@@ -112,7 +112,60 @@ const _defaultHandler = makeQueryHandler({
     element.querySelectorAll(selector),
 });
 
-const _builtInHandlers = new Map([['aria', ariaHandler]]);
+const pierceHandler = makeQueryHandler({
+  queryOne: (element, selector) => {
+    let found: Element | null = null;
+    const search = (root: Element | ShadowRoot) => {
+      const iter = document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT);
+      do {
+        const currentNode = iter.currentNode as HTMLElement;
+        if (currentNode.shadowRoot) {
+          search(currentNode.shadowRoot);
+        }
+        if (currentNode instanceof ShadowRoot) {
+          continue;
+        }
+        if (!found && currentNode.matches(selector)) {
+          found = currentNode;
+        }
+      } while (!found && iter.nextNode());
+    };
+    if (element instanceof Document) {
+      element = element.documentElement;
+    }
+    search(element);
+    return found;
+  },
+
+  queryAll: (element, selector) => {
+    const result: Element[] = [];
+    const collect = (root: Element | ShadowRoot) => {
+      const iter = document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT);
+      do {
+        const currentNode = iter.currentNode as HTMLElement;
+        if (currentNode.shadowRoot) {
+          collect(currentNode.shadowRoot);
+        }
+        if (currentNode instanceof ShadowRoot) {
+          continue;
+        }
+        if (currentNode.matches(selector)) {
+          result.push(currentNode);
+        }
+      } while (iter.nextNode());
+    };
+    if (element instanceof Document) {
+      element = element.documentElement;
+    }
+    collect(element);
+    return result;
+  },
+});
+
+const _builtInHandlers = new Map([
+  ['aria', ariaHandler],
+  ['pierce', pierceHandler],
+]);
 const _queryHandlers = new Map(_builtInHandlers);
 
 /**
diff --git a/test/queryselector.spec.ts b/test/queryselector.spec.ts
index 2b36c0cb..93a62a8f 100644
--- a/test/queryselector.spec.ts
+++ b/test/queryselector.spec.ts
@@ -68,6 +68,45 @@ describe('querySelector', function () {
     });
   });
 
+  describe('pierceHandler', function () {
+    beforeEach(async () => {
+      const { page } = getTestState();
+      await page.setContent(
+        `<script>
+        const div = document.createElement('div');
+        const shadowRoot = div.attachShadow({mode: 'open'});
+        const div1 = document.createElement('div');
+        div1.textContent = 'Hello';
+        div1.className = 'foo';
+        const div2 = document.createElement('div');
+        div2.textContent = 'World';
+        div2.className = 'foo';
+        shadowRoot.appendChild(div1);
+        shadowRoot.appendChild(div2);
+        document.documentElement.appendChild(div);
+        </script>`
+      );
+    });
+    it('should find first element in shadow', async () => {
+      const { page } = getTestState();
+      const div = await page.$('pierce/.foo');
+      const text = await div.evaluate(
+        (element: Element) => element.textContent
+      );
+      expect(text).toBe('Hello');
+    });
+    it('should find all elements in shadow', async () => {
+      const { page } = getTestState();
+      const divs = await page.$$('pierce/.foo');
+      const text = await Promise.all(
+        divs.map((div) =>
+          div.evaluate((element: Element) => element.textContent)
+        )
+      );
+      expect(text.join(' ')).toBe('Hello World');
+    });
+  });
+
   // The tests for $$eval are repeated later in this file in the test group 'QueryAll'.
   // This is done to also test a query handler where QueryAll returns an Element[]
   // as opposed to NodeListOf<Element>.