From 3ce02a95c8d8b3572edfa8089745b47fa8545cb6 Mon Sep 17 00:00:00 2001
From: Joel Einbinder <joel.einbinder@gmail.com>
Date: Thu, 25 Mar 2021 18:43:33 -0700
Subject: [PATCH] fix(selectors): properly generate selectors for tricky ids
 (#5940)

---
 src/server/supplements/injected/selectorGenerator.ts | 9 +++++++--
 test/selector-generator.spec.ts                      | 7 ++++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/server/supplements/injected/selectorGenerator.ts b/src/server/supplements/injected/selectorGenerator.ts
index 1d8a03ed27..deb32fdfe1 100644
--- a/src/server/supplements/injected/selectorGenerator.ts
+++ b/src/server/supplements/injected/selectorGenerator.ts
@@ -171,7 +171,8 @@ function buildCandidates(injectedScript: InjectedScript, element: Element): Sele
 
   const idAttr = element.getAttribute('id');
   if (idAttr && !isGuidLike(idAttr))
-    candidates.push({ engine: 'css', selector: `#${idAttr}`, score: 100 });
+    candidates.push({ engine: 'css', selector: makeSelectorForId(idAttr), score: 100 });
+
 
   candidates.push({ engine: 'css', selector: element.nodeName.toLocaleLowerCase(), score: 200 });
   return candidates;
@@ -209,6 +210,10 @@ function parentElementOrShadowHost(element: Element): Element | null {
   return null;
 }
 
+function makeSelectorForId(id: string) {
+  return /^[a-zA-Z][a-zA-Z0-9\-\_]+$/.test(id) ? '#' + id : `[id="${id}"]`;
+}
+
 function cssFallback(injectedScript: InjectedScript, targetElement: Element): SelectorToken {
   const kFallbackScore = 10000000;
   const root: Node = targetElement.ownerDocument;
@@ -230,7 +235,7 @@ function cssFallback(injectedScript: InjectedScript, targetElement: Element): Se
     // Element ID is the strongest signal, use it.
     let bestTokenForLevel: string = '';
     if (element.id) {
-      const token = /^[a-zA-Z][a-zA-Z0-9\-\_]+$/.test(element.id) ? '#' + element.id : `[id="${element.id}"]`;
+      const token = makeSelectorForId(element.id);
       const selector = uniqueCSSSelector(token);
       if (selector)
         return { engine: 'css', selector, score: kFallbackScore };
diff --git a/test/selector-generator.spec.ts b/test/selector-generator.spec.ts
index 8ae9577a79..d778ddb5a4 100644
--- a/test/selector-generator.spec.ts
+++ b/test/selector-generator.spec.ts
@@ -260,7 +260,7 @@ describe('selector generator', (suite, { mode }) => {
     const [frame] = await Promise.all([
       page.waitForEvent('frameattached'),
       page.evaluate(() => {
-        return new Promise(f => {
+        return new Promise<void>(f => {
           const iframe = document.createElement('iframe');
           iframe.onload = () => {
             iframe.contentDocument.body.innerHTML = '<div>Target</div>';
@@ -279,4 +279,9 @@ describe('selector generator', (suite, { mode }) => {
       expect(await generate(page, '[name=bar]')).toBe(`${tagName}[name="bar"]`);
     }
   });
+
+  it('should work with tricky ids', async ({page}) => {
+    await page.setContent(`<button id="this:is-my-tricky.id"><span></span></button>`);
+    expect(await generate(page, 'button')).toBe('[id="this:is-my-tricky.id"]');
+  });
 });