feat(browser): custom locators API (#7993)

Author: sheremet-va

docs/guide/browser/locators.md

@@ -955,3 +955,76 @@ test('works correctly', async () => {
 })
 ```
 :::
+
+## Custom Locators <Version>3.2.0</Version> <Badge type="danger">advanced</Badge> {#custom-locators}
+
+You can extend built-in locators API by defining an object of locator factories. These methods will exist as methods on the `page` object and any created locator.
+
+These locators can be useful if built-in locators are not enough. For example, when you use a custom framework for your UI.
+
+The locator factory needs to return a selector string or a locator itself.
+
+::: tip
+The selector syntax is identical to Playwright locators. Please, read [their guide](https://playwright.dev/docs/other-locators) to better understand how to work with them.
+:::
+
+```ts
+import { locators } from '@vitest/browser/context'
+
+locators.extend({
+  getByArticleTitle(title) {
+    return `[data-title="${title}"]`
+  },
+  getByArticleCommentsCount(count) {
+    return `.comments :text("${count} comments")`
+  },
+  async previewComments() {
+    // you have access to the current locator via "this"
+    // beware that if the method was called on `page`, `this` will be `page`,
+    // not the locator!
+    if (this !== page) {
+      await this.click()
+    }
+    // ...
+  }
+})
+
+// if you are using typescript, you can extend LocatorSelectors interface
+// to have the autocompletion in locators.extend, page.* and locator.* methods
+declare module '@vitest/browser/context' {
+  interface LocatorSelectors {
+    // if the custom method returns a string, it will be converted into a locator
+    // if it returns anything else, then it will be returned as usual
+    getByArticleTitle(title: string): Locator
+    getByArticleCommentsCount(count: number): Locator
+
+    // Vitest will return a promise and won't try to convert it into a locator
+    previewComments(this: Locator): Promise<void>
+  }
+}
+```
+
+If the method is called on the global `page` object, then selector will be applied to the whole page. In the example bellow, `getByArticleTitle` will find all elements with an attribute `data-title` with the value of `title`. However, if the method is called on the locator, then it will be scoped to that locator.
+
+```html
+<article data-title="Hello, World!">
+  Hello, World!
+  <button id="comments">2 comments</button>
+</article>
+
+<article data-title="Hello, Vitest!">
+  Hello, Vitest!
+  <button id="comments">0 comments</button>
+</article>
+```
+
+```ts
+const articles = page.getByRole('article')
+const worldArticle = page.getByArticleTitle('Hello, World!') // ✅
+const commentsElement = worldArticle.getByArticleCommentsCount(2) // ✅
+const wrongCommentsElement = worldArticle.getByArticleCommentsCount(0) // ❌
+const wrongElement = page.getByArticleTitle('No Article!') // ❌
+
+await commentsElement.previewComments() // ✅
+await wrongCommentsElement.previewComments() // ❌
+```

packages/browser/context.d.ts

@@ -580,5 +580,14 @@ export interface BrowserPage extends LocatorSelectors {
   elementLocator(element: Element): Locator
 }
 
+export interface BrowserLocators {
+  createElementLocators(element: Element): LocatorSelectors
+  extend(methods: {
+    [K in keyof LocatorSelectors]?: (...args: Parameters<LocatorSelectors[K]>) => ReturnType<LocatorSelectors[K]> | string
+  }): void
+}
+
+export const locators: BrowserLocators
+
 export const page: BrowserPage
 export const cdp: () => CDPSession

packages/browser/context.js

@@ -6,6 +6,7 @@ export const server = null
 export const userEvent = null
 export const cdp = null
 export const commands = null
+export const locators = null
 
 const pool = globalThis.__vitest_worker__?.ctx?.pool
 

packages/browser/src/client/tester/context.ts

@@ -1,12 +1,15 @@
 import type { Options as TestingLibraryOptions, UserEvent as TestingLibraryUserEvent } from '@testing-library/user-event'
 import type { RunnerTask } from 'vitest'
 import type {
+  BrowserLocators,
   BrowserPage,
   Locator,
   UserEvent,
 } from '../../../context'
 import type { IframeViewportEvent } from '../client'
 import type { BrowserRunnerState } from '../utils'
+import type { Locator as LocatorAPI } from './locators/index'
+import { getElementLocatorSelectors } from '@vitest/browser/utils'
 import { ensureAwaited, getBrowserState, getWorkerState } from '../utils'
 import { convertElementToCssSelector, processTimeoutOptions } from './utils'
 
@@ -317,9 +320,12 @@ export const page: BrowserPage = {
   elementLocator() {
     throw new Error(`Method "elementLocator" is not implemented in the "${provider}" provider.`)
   },
+  _createLocator() {
+    throw new Error(`Method "_createLocator" is not implemented in the "${provider}" provider.`)
+  },
   extend(methods) {
     for (const key in methods) {
-      (page as any)[key] = (methods as any)[key]
+      (page as any)[key] = (methods as any)[key].bind(page)
     }
     return page
   },
@@ -348,3 +354,35 @@ function convertToSelector(elementOrLocator: Element | Locator): string {
 function getTaskFullName(task: RunnerTask): string {
   return task.suite ? `${getTaskFullName(task.suite)} ${task.name}` : task.name
 }
+
+export const locators: BrowserLocators = {
+  createElementLocators: getElementLocatorSelectors,
+  extend(methods) {
+    const Locator = page._createLocator('css=body').constructor as typeof LocatorAPI
+    for (const method in methods) {
+      const cb = (methods as any)[method] as (...args: any[]) => string | Locator
+      // @ts-expect-error types are hard to make work
+      Locator.prototype[method] = function (...args: any[]) {
+        const selectorOrLocator = cb.call(this, ...args)
+        if (typeof selectorOrLocator === 'string') {
+          return this.locator(selectorOrLocator)
+        }
+        return selectorOrLocator
+      }
+      page[method as 'getByRole'] = function (...args: any[]) {
+        const selectorOrLocator = cb.call(this, ...args)
+        if (typeof selectorOrLocator === 'string') {
+          return page._createLocator(selectorOrLocator)
+        }
+        return selectorOrLocator
+      }
+    }
+  },
+}
+
+declare module '@vitest/browser/context' {
+  interface BrowserPage {
+    /** @internal */
+    _createLocator: (selector: string) => Locator
+  }
+}

packages/browser/src/client/tester/locators/playwright.ts

@@ -35,6 +35,9 @@ page.extend({
     return new PlaywrightLocator(getByTitleSelector(title, options))
   },
 
+  _createLocator(selector: string) {
+    return new PlaywrightLocator(selector)
+  },
   elementLocator(element: Element) {
     return new PlaywrightLocator(
       selectorEngine.generateSelectorSimple(element),

packages/browser/src/client/tester/locators/preview.ts

@@ -35,6 +35,9 @@ page.extend({
     return new PreviewLocator(getByTitleSelector(title, options))
   },
 
+  _createLocator(selector: string) {
+    return new PreviewLocator(selector)
+  },
   elementLocator(element: Element) {
     return new PreviewLocator(
       selectorEngine.generateSelectorSimple(element),

packages/browser/src/client/tester/locators/webdriverio.ts

@@ -37,6 +37,9 @@ page.extend({
     return new WebdriverIOLocator(getByTitleSelector(title, options))
   },
 
+  _createLocator(selector: string) {
+    return new WebdriverIOLocator(selector)
+  },
   elementLocator(element: Element) {
     return new WebdriverIOLocator(selectorEngine.generateSelectorSimple(element))
   },

packages/browser/src/node/plugins/pluginContext.ts

@@ -49,7 +49,7 @@ async function generateContextFile(
   const distContextPath = slash(`/@fs/${resolve(__dirname, 'context.js')}`)
 
   return `
-import { page, createUserEvent, cdp } from '${distContextPath}'
+import { page, createUserEvent, cdp, locators } from '${distContextPath}'
 ${userEventNonProviderImport}
 
 export const server = {
@@ -64,7 +64,7 @@ export const server = {
 }
 export const commands = server.commands
 export const userEvent = createUserEvent(_userEventSetup)
-export { page, cdp }
+export { page, cdp, locators }
 `
 }
 

test/browser/fixtures/locators-custom/basic.test.tsx

@@ -0,0 +1,86 @@
+import { type Locator, locators, page } from '@vitest/browser/context';
+import { beforeEach, expect, test } from 'vitest';
+
+declare module '@vitest/browser/context' {
+  interface LocatorSelectors {
+    getByCustomTitle: (title: string) => Locator
+    getByNestedTitle: (title: string) => Locator
+    updateHtml(this: Locator, html: string): Promise<void>
+    updateDocumentHtml(this: BrowserPage, html: string): Promise<void>
+  }
+}
+
+locators.extend({
+  getByCustomTitle(title) {
+    return `[data-title="${title}"]`
+  },
+  getByNestedTitle(title) {
+    return `[data-parent] >> [data-title="${title}"]`
+  },
+  async updateHtml(this: Locator, html) {
+    this.element().innerHTML = html
+  },
+  async updateDocumentHtml(html) {
+    document.body.innerHTML = html
+  },
+})
+
+beforeEach(() => {
+  document.body.innerHTML = ''
+})
+
+test('new selector works on both page and locator', async () => {
+  document.body.innerHTML = `
+  <article>
+    <h1>Hello World</h1>
+    <div data-title="Hello World">Text Content</div>
+  </article>
+  `
+
+  await expect.element(page.getByCustomTitle('Hello World')).toBeVisible()
+  await expect.element(page.getByRole('article').getByCustomTitle('Hello World')).toBeVisible()
+
+  await expect.element(page.getByCustomTitle('NonExisting Title')).not.toBeInTheDocument()
+})
+
+test('new nested selector works on both page and locator', async () => {
+  document.body.innerHTML = `
+  <article>
+    <h1>Hello World</h1>
+    <div data-parent>
+      <div data-title="Hello World">Text Content</div>
+    </div>
+  </article>
+  `
+
+  await expect.element(page.getByNestedTitle('Hello World')).toBeVisible()
+  await expect.element(page.getByRole('article').getByNestedTitle('Hello World')).toBeVisible()
+
+  await expect.element(page.getByNestedTitle('NonExisting Title')).not.toBeInTheDocument()
+})
+
+test('new added method works on the locator', async () => {
+  document.body.innerHTML = `
+    <div data-title="Hello World">Text Content</div>
+  `
+
+  const title = page.getByCustomTitle('Hello World')
+
+  await expect.element(title).toHaveTextContent('Text Content')
+
+  await title.updateHtml('New Content')
+
+  await expect.element(title).toHaveTextContent('New Content')
+})
+
+test('new added method works on the page', async () => {
+  document.body.innerHTML = `
+    Hello World
+  `
+
+  expect(document.body).toHaveTextContent('Hello World')
+
+  await page.updateDocumentHtml('New Content')
+
+  expect(document.body).toHaveTextContent('New Content')
+})

test/browser/fixtures/locators-custom/vitest.config.ts

@@ -0,0 +1,15 @@
+import { fileURLToPath } from 'node:url'
+import { defineConfig } from 'vitest/config'
+import { instances, provider } from '../../settings'
+
+export default defineConfig({
+  cacheDir: fileURLToPath(new URL("./node_modules/.vite", import.meta.url)),
+  test: {
+    browser: {
+      enabled: true,
+      provider,
+      headless: true,
+      instances,
+    },
+  },
+})