feat(fixture): support chaining locator queries with locator.within()

// Synchronous
```ts
test('chaining synchronous queries', async ({screen}) => {
  const locator = screen.getByRole('figure').within().getByText('Some image')

  expect(await locator.textContent()).toEqual('Some image')
})
```

// Synchronous + Asynchronous
```ts
test('chaining multiple asynchronous queries between synchronous queries', async ({screen}) => {
  const locator = await screen
    .getByTestId('modal-container')
    .within()
    .findByRole('dialog')
    .within()
    .findByRole('alert')
    .within()
    .getByRole('button', {name: 'Close'})

  expect(await locator.textContent()).toEqual('Close')
})
```

Author: jrolfs

lib/fixture/locator/fixtures.ts

@@ -1,10 +1,11 @@
-import type {Locator, PlaywrightTestArgs, TestFixture} from '@playwright/test'
-import {Page, selectors} from '@playwright/test'
+import type {PlaywrightTestArgs, TestFixture} from '@playwright/test'
+import {selectors} from '@playwright/test'
 
 import type {TestingLibraryDeserializedFunction as DeserializedFunction} from '../helpers'
 import type {
   Config,
   LocatorQueries as Queries,
+  QueryRoot,
   Screen,
   SelectorEngine,
   SynchronousQuery,
@@ -47,10 +48,10 @@ const withinFixture: TestFixture<Within, TestArguments> = async (
   {asyncUtilExpectedState, asyncUtilTimeout},
   use,
 ) =>
-  use(<Root extends Page | Locator>(root: Root) =>
+  use(<Root extends QueryRoot>(root: Root) =>
     'goto' in root
-      ? screenFor(root, {asyncUtilExpectedState, asyncUtilTimeout}).proxy
-      : (queriesFor(root, {asyncUtilExpectedState, asyncUtilTimeout}) as WithinReturn<Root>),
+      ? (screenFor(root, {asyncUtilExpectedState, asyncUtilTimeout}).proxy as WithinReturn<Root>)
+      : (queriesFor<Root>(root, {asyncUtilExpectedState, asyncUtilTimeout}) as WithinReturn<Root>),
   )
 
 type SynchronousQueryParameters = Parameters<Queries[SynchronousQuery]>

lib/fixture/locator/index.ts

@@ -1,3 +1,6 @@
+export type {Queries} from './fixtures'
+export type {LocatorPromise} from './queries'
+
 export {
   installTestingLibraryFixture,
   options,
@@ -6,5 +9,4 @@ export {
   screenFixture,
   withinFixture,
 } from './fixtures'
-export type {Queries} from './fixtures'
 export {queriesFor} from './queries'

lib/fixture/locator/queries.ts

@@ -1,5 +1,5 @@
-import type {Locator, Page} from '@playwright/test'
-import {errors} from '@playwright/test'
+import type {Page} from '@playwright/test'
+import {Locator, errors} from '@playwright/test'
 import {queries} from '@testing-library/dom'
 
 import {replacer} from '../helpers'
@@ -9,14 +9,19 @@ import type {
   FindQuery,
   GetQuery,
   LocatorQueries as Queries,
+  QueriesReturn,
   Query,
   QueryQuery,
+  QueryRoot,
   Screen,
   SynchronousQuery,
+  TestingLibraryLocator,
 } from '../types'
 
 import {includes, queryToSelector} from './helpers'
 
+type SynchronousQueryParameters = Parameters<Queries[SynchronousQuery]>
+
 const isAllQuery = (query: Query): query is AllQuery => query.includes('All')
 
 const isFindQuery = (query: Query): query is FindQuery => query.startsWith('find')
@@ -29,60 +34,115 @@ const synchronousQueryNames = allQueryNames.filter(isNotFindQuery)
 const findQueryToGetQuery = (query: FindQuery) => query.replace(/^find/, 'get') as GetQuery
 const findQueryToQueryQuery = (query: FindQuery) => query.replace(/^find/, 'query') as QueryQuery
 
-const createFindQuery =
-  (
-    pageOrLocator: Page | Locator,
-    query: FindQuery,
-    {asyncUtilTimeout, asyncUtilExpectedState}: Partial<Config> = {},
-  ) =>
-  async (...[id, options, waitForElementOptions]: Parameters<Queries[FindQuery]>) => {
-    const synchronousOptions = ([id, options] as const).filter(Boolean)
-
-    const locator = pageOrLocator.locator(
-      `${queryToSelector(findQueryToQueryQuery(query))}=${JSON.stringify(
-        synchronousOptions,
-        replacer,
-      )}`,
-    )
-
-    const {state: expectedState = asyncUtilExpectedState, timeout = asyncUtilTimeout} =
-      waitForElementOptions ?? {}
-
-    try {
-      await locator.first().waitFor({state: expectedState, timeout})
-    } catch (error) {
-      // In the case of a `waitFor` timeout from Playwright, we want to
-      // surface the appropriate error from Testing Library, so run the
-      // query one more time as `get*` knowing that it will fail with the
-      // error that we want the user to see instead of the `TimeoutError`
-      if (error instanceof errors.TimeoutError) {
-        const timeoutLocator = pageOrLocator
-          .locator(
-            `${queryToSelector(findQueryToGetQuery(query))}=${JSON.stringify(
-              synchronousOptions,
-              replacer,
-            )}`,
-          )
-          .first()
-
-        // Handle case where element is attached, but hidden, and the expected
-        // state is set to `visible`. In this case, dereferencing the
-        // `Locator` instance won't throw a `get*` query error, so just
-        // surface the original Playwright timeout error
-        if (expectedState === 'visible' && !(await timeoutLocator.isVisible())) {
-          throw error
-        }
+class LocatorPromise extends Promise<Locator> {
+  /**
+   * Wrap an `async` function `Promise` return value in a `LocatorPromise`.
+   * This allows us to use `async/await` and still return a custom
+   * `LocatorPromise` instance instead of `Promise`.
+   *
+   * @param fn
+   * @returns
+   */
+  static wrap<A extends any[]>(fn: (...args: A) => Promise<Locator>, config: Partial<Config>) {
+    return (...args: A) => LocatorPromise.from(fn(...args), config)
+  }
 
-        // In all other cases, dereferencing the `Locator` instance here should
-        // cause the above `get*` query to throw an error in Testing Library
-        return timeoutLocator.waitFor({state: expectedState, timeout})
-      }
+  static from(promise: Promise<Locator>, config: Partial<Config>) {
+    return new LocatorPromise((resolve, reject) => {
+      promise.then(resolve).catch(reject)
+    }, config)
+  }
+
+  config: Partial<Config>
 
-      throw error
-    }
+  constructor(
+    executor: (
+      resolve: (value: Locator | PromiseLike<Locator>) => void,
+      reject: (reason?: any) => void,
+    ) => void,
+    config: Partial<Config>,
+  ) {
+    super(executor)
+
+    this.config = config
+  }
 
-    return locator
+  within() {
+    // eslint-disable-next-line @typescript-eslint/no-use-before-define
+    return queriesFor(this, this.config)
   }
+}
+
+const locatorFor = (
+  root: Exclude<QueryRoot, Promise<any>>,
+  query: SynchronousQuery,
+  options: SynchronousQueryParameters,
+) => root.locator(`${queryToSelector(query)}=${JSON.stringify(options, replacer)}`)
+
+const augmentedLocatorFor = (
+  root: Exclude<QueryRoot, Promise<any>>,
+  query: SynchronousQuery,
+  options: SynchronousQueryParameters,
+  config: Partial<Config>,
+) => {
+  const locator = locatorFor(root, query, options)
+
+  return new Proxy(locator, {
+    get(target, property, receiver) {
+      return property === 'within'
+        ? // eslint-disable-next-line @typescript-eslint/no-use-before-define
+          () => queriesFor(target, config)
+        : Reflect.get(target, property, receiver)
+    },
+  }) as TestingLibraryLocator
+}
+
+const createFindQuery = (
+  root: QueryRoot,
+  query: FindQuery,
+  {asyncUtilTimeout, asyncUtilExpectedState}: Partial<Config> = {},
+) =>
+  LocatorPromise.wrap(
+    async (...[id, options, waitForElementOptions]: Parameters<Queries[FindQuery]>) => {
+      const settledRoot = root instanceof LocatorPromise ? await root : root
+      const synchronousOptions = (options ? [id, options] : [id]) as SynchronousQueryParameters
+
+      const locator = locatorFor(settledRoot, findQueryToQueryQuery(query), synchronousOptions)
+      const {state: expectedState = asyncUtilExpectedState, timeout = asyncUtilTimeout} =
+        waitForElementOptions ?? {}
+
+      try {
+        await locator.first().waitFor({state: expectedState, timeout})
+      } catch (error) {
+        // In the case of a `waitFor` timeout from Playwright, we want to
+        // surface the appropriate error from Testing Library, so run the
+        // query one more time as `get*` knowing that it will fail with the
+        // error that we want the user to see instead of the `TimeoutError`
+        if (error instanceof errors.TimeoutError) {
+          const timeoutLocator = locatorFor(
+            settledRoot,
+            findQueryToGetQuery(query),
+            synchronousOptions,
+          ).first()
+
+          // Handle case where element is attached, but hidden, and the expected
+          // state is set to `visible`. In this case, dereferencing the
+          // `Locator` instance won't throw a `get*` query error, so just
+          // surface the original Playwright timeout error
+          if (expectedState === 'visible' && !(await timeoutLocator.isVisible())) {
+            throw error
+          }
+
+          // In all other cases, dereferencing the `Locator` instance here should
+          // cause the above `get*` query to throw an error in Testing Library
+          await timeoutLocator.waitFor({state: expectedState, timeout})
+        }
+      }
+
+      return locator
+    },
+    {asyncUtilExpectedState, asyncUtilTimeout},
+  )
 
 /**
  * Given a `Page` or `Locator` instance, return an object of Testing Library
@@ -93,21 +153,26 @@ const createFindQuery =
  * should use the `locatorFixtures` with **@playwright/test** instead.
  * @see {@link locatorFixtures}
  *
- * @param pageOrLocator `Page` or `Locator` instance to use as the query root
+ * @param root `Page` or `Locator` instance to use as the query root
  * @param config Testing Library configuration to apply to queries
  *
  * @returns object containing scoped Testing Library query methods
  */
-const queriesFor = (pageOrLocator: Page | Locator, config?: Partial<Config>) =>
+const queriesFor = <Root extends QueryRoot>(
+  root: Root,
+  config: Partial<Config>,
+): QueriesReturn<Root> =>
   allQueryNames.reduce(
     (rest, query) => ({
       ...rest,
       [query]: isFindQuery(query)
-        ? createFindQuery(pageOrLocator, query, config)
-        : (...args: Parameters<Queries[SynchronousQuery]>) =>
-            pageOrLocator.locator(`${queryToSelector(query)}=${JSON.stringify(args, replacer)}`),
+        ? createFindQuery(root, query, config)
+        : (...options: SynchronousQueryParameters) =>
+            root instanceof LocatorPromise
+              ? root.then(r => locatorFor(r, query, options))
+              : augmentedLocatorFor(root, query, options, config),
     }),
-    {} as Queries,
+    {} as QueriesReturn<Root>,
   )
 
 const screenFor = (page: Page, config: Partial<Config>) =>
@@ -119,4 +184,12 @@ const screenFor = (page: Page, config: Partial<Config>) =>
     },
   }) as {proxy: Screen; revoke: () => void}
 
-export {allQueryNames, isAllQuery, isNotFindQuery, queriesFor, screenFor, synchronousQueryNames}
+export {
+  LocatorPromise,
+  allQueryNames,
+  isAllQuery,
+  isNotFindQuery,
+  queriesFor,
+  screenFor,
+  synchronousQueryNames,
+}

lib/fixture/types.ts

@@ -5,6 +5,7 @@ import {queries} from '@testing-library/dom'
 import type {Config as CommonConfig} from '../common'
 
 import {reviver} from './helpers'
+import type {LocatorPromise} from './locator'
 
 /**
  * This type was copied across from Playwright
@@ -22,15 +23,23 @@ export type SelectorEngine = {
   queryAll(root: HTMLElement, selector: string): HTMLElement[]
 }
 
+type KebabCase<S> = S extends `${infer C}${infer T}`
+  ? T extends Uncapitalize<T>
+    ? `${Uncapitalize<C>}${KebabCase<T>}`
+    : `${Uncapitalize<C>}-${KebabCase<T>}`
+  : S
+
 type Queries = typeof queries
 type WaitForState = Exclude<Parameters<Locator['waitFor']>[0], undefined>['state']
 type AsyncUtilExpectedState = Extract<WaitForState, 'visible' | 'attached'>
 
+export type TestingLibraryLocator = Locator & {within: () => LocatorQueries}
+
 type ConvertQuery<Query extends Queries[keyof Queries]> = Query extends (
   el: HTMLElement,
   ...rest: infer Rest
 ) => HTMLElement | (HTMLElement[] | null) | (HTMLElement | null)
-  ? (...args: Rest) => Locator
+  ? (...args: Rest) => TestingLibraryLocator
   : Query extends (
       el: HTMLElement,
       id: infer Id,
@@ -41,23 +50,31 @@ type ConvertQuery<Query extends Queries[keyof Queries]> = Query extends (
       id: Id,
       options?: Options,
       waitForOptions?: WaitForOptions & {state?: AsyncUtilExpectedState},
-    ) => Promise<Locator>
+    ) => LocatorPromise
   : never
 
-type KebabCase<S> = S extends `${infer C}${infer T}`
-  ? T extends Uncapitalize<T>
-    ? `${Uncapitalize<C>}${KebabCase<T>}`
-    : `${Uncapitalize<C>}-${KebabCase<T>}`
-  : S
-
 export type LocatorQueries = {[K in keyof Queries]: ConvertQuery<Queries[K]>}
 
-export type WithinReturn<Root extends Locator | Page> = Root extends Page ? Screen : LocatorQueries
+type ConvertQueryDeferred<Query extends LocatorQueries[keyof LocatorQueries]> = Query extends (
+  ...rest: infer Rest
+) => any
+  ? (...args: Rest) => LocatorPromise
+  : never
+
+export type DeferredLocatorQueries = {
+  [K in keyof LocatorQueries]: ConvertQueryDeferred<LocatorQueries[K]>
+}
+
+export type WithinReturn<Root extends QueryRoot> = Root extends Page ? Screen : QueriesReturn<Root>
+export type QueriesReturn<Root extends QueryRoot> = Root extends LocatorPromise
+  ? DeferredLocatorQueries
+  : LocatorQueries
+
+export type QueryRoot = Page | Locator | LocatorPromise
 export type Screen = LocatorQueries & Page
-export type Within = <Root extends Locator | Page>(locator: Root) => WithinReturn<Root>
+export type Within = <Root extends QueryRoot>(locator: Root) => WithinReturn<Root>
 
 export type Query = keyof Queries
-
 export type AllQuery = Extract<Query, `${string}All${string}`>
 export type FindQuery = Extract<Query, `find${string}`>
 export type GetQuery = Extract<Query, `get${string}`>

playwright.config.ts

@@ -4,6 +4,7 @@ const config: PlaywrightTestConfig = {
   reporter: 'list',
   testDir: 'test/fixture',
   use: {actionTimeout: 3000},
+  timeout: 5 * 1000,
 }
 
 export default config