diff --git a/packages/playwright-core/src/client/waiter.ts b/packages/playwright-core/src/client/waiter.ts index 1b3ffbe78d..7c239baa2f 100644 --- a/packages/playwright-core/src/client/waiter.ts +++ b/packages/playwright-core/src/client/waiter.ts @@ -17,7 +17,8 @@ import type { EventEmitter } from 'events'; import { rewriteErrorMessage } from '../utils/stackTrace'; import { TimeoutError } from './errors'; -import { createGuid } from '../utils'; +import { createGuid, zones } from '../utils'; +import type { Zone } from '../utils'; import type * as channels from '@protocol/channels'; import type { ChannelOwner } from './channelOwner'; @@ -29,10 +30,13 @@ export class Waiter { private _channelOwner: ChannelOwner; private _waitId: string; private _error: string | undefined; + private _savedZone: Zone | undefined; constructor(channelOwner: ChannelOwner, event: string) { this._waitId = createGuid(); this._channelOwner = channelOwner; + this._savedZone = zones.currentZone(); + this._channelOwner._channel.waitForEventInfo({ info: { waitId: this._waitId, phase: 'before', event } }).catch(() => {}); this._dispose = [ () => this._channelOwner._wrapApiCall(async () => { @@ -46,12 +50,12 @@ export class Waiter { } async waitForEvent(emitter: EventEmitter, event: string, predicate?: (arg: T) => boolean | Promise): Promise { - const { promise, dispose } = waitForEvent(emitter, event, predicate); + const { promise, dispose } = waitForEvent(emitter, event, this._savedZone, predicate); return await this.waitForPromise(promise, dispose); } rejectOnEvent(emitter: EventEmitter, event: string, error: Error | (() => Error), predicate?: (arg: T) => boolean | Promise) { - const { promise, dispose } = waitForEvent(emitter, event, predicate); + const { promise, dispose } = waitForEvent(emitter, event, this._savedZone, predicate); this._rejectOn(promise.then(() => { throw (typeof error === 'function' ? error() : error); }), dispose); } @@ -103,19 +107,22 @@ export class Waiter { } } -function waitForEvent(emitter: EventEmitter, event: string, predicate?: (arg: T) => boolean | Promise): { promise: Promise, dispose: () => void } { +function waitForEvent(emitter: EventEmitter, event: string, savedZone: Zone | undefined, predicate?: (arg: T) => boolean | Promise): { promise: Promise, dispose: () => void } { let listener: (eventArg: any) => void; const promise = new Promise((resolve, reject) => { listener = async (eventArg: any) => { - try { - if (predicate && !(await predicate(eventArg))) - return; - emitter.removeListener(event, listener); - resolve(eventArg); - } catch (e) { - emitter.removeListener(event, listener); - reject(e); - } + // Reset apiZone and expectZone, but restore step data. + await zones.runInZone(savedZone?.copyWithoutTypes(['apiZone', 'expectZone']), async () => { + try { + if (predicate && !(await predicate(eventArg))) + return; + emitter.removeListener(event, listener); + resolve(eventArg); + } catch (e) { + emitter.removeListener(event, listener); + reject(e); + } + }); }; emitter.addListener(event, listener); }); diff --git a/packages/playwright-core/src/utils/zones.ts b/packages/playwright-core/src/utils/zones.ts index e6fabac0f1..8dfa6404fc 100644 --- a/packages/playwright-core/src/utils/zones.ts +++ b/packages/playwright-core/src/utils/zones.ts @@ -19,49 +19,52 @@ import { AsyncLocalStorage } from 'async_hooks'; export type ZoneType = 'apiZone' | 'expectZone' | 'stepZone'; class ZoneManager { - private readonly _asyncLocalStorage = new AsyncLocalStorage|undefined>(); + private readonly _asyncLocalStorage = new AsyncLocalStorage(); run(type: ZoneType, data: T, func: () => R): R { - const previous = this._asyncLocalStorage.getStore(); - const zone = new Zone(previous, type, data); + const current = this._asyncLocalStorage.getStore(); + const zone = Zone.createWithData(current, type, data); + return this.runInZone(zone, func); + } + + runInZone(zone: Zone | undefined, func: () => R): R { return this._asyncLocalStorage.run(zone, func); } zoneData(type: ZoneType): T | undefined { - for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) { - if (zone.type === type) - return zone.data as T; - } - return undefined; + const zone = this._asyncLocalStorage.getStore(); + return zone?.get(type); + } + + currentZone(): Zone | undefined { + return this._asyncLocalStorage.getStore(); } exitZones(func: () => R): R { return this._asyncLocalStorage.run(undefined, func); } - - printZones() { - const zones = []; - for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) { - let str = zone.type; - if (zone.type === 'apiZone') - str += `(${(zone.data as any).apiName})`; - zones.push(str); - - } - // eslint-disable-next-line no-console - console.log('zones: ', zones.join(' -> ')); - } } -class Zone { - readonly type: ZoneType; - readonly data: T; - readonly previous: Zone | undefined; +export class Zone { + private readonly store: Map; - constructor(previous: Zone | undefined, type: ZoneType, data: T) { - this.type = type; - this.data = data; - this.previous = previous; + static createWithData(currentZone: Zone | undefined, type: ZoneType, data: unknown) { + const store = new Map(currentZone?.store.entries() ?? []); + store.set(type, data); + return new Zone(store); + } + + private constructor(store: Map) { + this.store = store; + } + + copyWithoutTypes(types: ZoneType[]): Zone { + const store = new Map(this.store.entries().filter(([type]) => !types.includes(type))); + return new Zone(store); + } + + get(type: ZoneType): T | undefined { + return this.store.get(type) as T | undefined; } }