fix: disable restored mocks (#10413)

* fix: disable mocks after restore

* fix: mock response error

packages/webdriverio/src/commands/browser/mock.ts

@@ -4,8 +4,10 @@ import WebDriverNetworkInterception from '../../utils/interception/webdriver.js'
 import { getBrowserObject } from '../../utils/index.js'
 import type { Mock } from '../../types.js'
 import type { MockFilterOptions } from '../../utils/interception/types.js'
+import type { CDPSession } from 'puppeteer-core/lib/esm/puppeteer/common/Connection.js'
 
 export const SESSION_MOCKS: Record<string, Set<Interception>> = {}
+export const CDP_SESSIONS: Record<string, CDPSession> = {}
 
 /**
  * Mock the response of a request. You can define a mock based on a matching
@@ -158,7 +160,7 @@ export async function mock (
             page = pages[0]
         }
 
-        const client = await page.target().createCDPSession()
+        const client = CDP_SESSIONS[handle] = await page.target().createCDPSession()
         await client.send('Fetch.enable', {
             patterns: [{ requestStage: 'Request' }, { requestStage: 'Response' }]
         })

packages/webdriverio/src/commands/browser/mockRestoreAll.ts

@@ -32,7 +32,7 @@ export async function mockRestoreAll () {
     for (const [handle, mocks] of Object.entries(SESSION_MOCKS)) {
         log.trace(`Clearing mocks for ${handle}`)
         for (const mock of mocks) {
-            mock.restore()
+            await mock.restore()
         }
     }
 }

packages/webdriverio/src/commands/mock/restore.ts

@@ -1,5 +1,6 @@
 /**
  * Does everything that `mock.clear()` does, and also removes any mocked return values or implementations.
+ * Restored mock does not emit events and could not mock responses.
  *
  * <example>
     :addValue.js

packages/webdriverio/src/types.ts

@@ -116,7 +116,7 @@ type MultiRemoteElementCommands = {
 }
 
 export type MultiRemoteBrowserCommandsType = {
-    [K in keyof Omit<BrowserCommandsType, ElementCommandNames | 'SESSION_MOCKS'>]: (...args: Parameters<BrowserCommandsType[K]>) => Promise<ThenArg<ReturnType<BrowserCommandsType[K]>>[]>
+    [K in keyof Omit<BrowserCommandsType, ElementCommandNames | 'SESSION_MOCKS' | 'CDP_SESSIONS'>]: (...args: Parameters<BrowserCommandsType[K]>) => Promise<ThenArg<ReturnType<BrowserCommandsType[K]>>[]>
 } & MultiRemoteElementCommands
 export type MultiRemoteElementCommandsType = {
     [K in keyof Omit<ElementCommandsType, ElementCommandNames>]: (...args: Parameters<ElementCommandsType[K]>) => Promise<ThenArg<ReturnType<ElementCommandsType[K]>>[]>

packages/webdriverio/src/utils/index.ts

@@ -23,7 +23,7 @@ import type { CustomStrategyReference } from '../types.js'
 
 const log = logger('webdriverio')
 const INVALID_SELECTOR_ERROR = 'selector needs to be typeof `string` or `function`'
-const IGNORED_COMMAND_FILE_EXPORTS = ['SESSION_MOCKS']
+const IGNORED_COMMAND_FILE_EXPORTS = ['SESSION_MOCKS', 'CDP_SESSIONS']
 
 declare global {
     interface Window { __wdio_element: Record<string, HTMLElement> }

packages/webdriverio/src/utils/interception/devtools.ts

@@ -8,6 +8,7 @@ import Interception from './index.js'
 import type { Matches, MockOverwrite, MockResponseParams } from './types.js'
 import { containsHeaderObject } from '../index.js'
 import { ERROR_REASON } from '../../constants.js'
+import { CDP_SESSIONS, SESSION_MOCKS } from '../../commands/browser/mock.js'
 
 const log = logger('webdriverio')
 
@@ -26,18 +27,21 @@ type Event = {
     request: Matches & { mockedResponse: string | Buffer }
     responseStatusCode?: number
     responseHeaders: HeaderEntry[]
+    responseErrorReason?: string
 }
 
 type ExpectParameter<T> = ((param: T) => boolean) | T;
 
 export default class DevtoolsInterception extends Interception {
+    private restored = false
+
     static handleRequestInterception (client: CDPSession, mocks: Set<Interception>): (event: Event) => Promise<void | ClientResponse> {
         return async (event) => {
             // responseHeaders and responseStatusCode are only present in Response stage
             // https://chromedevtools.github.io/devtools-protocol/tot/Fetch/#event-requestPaused
-            const isRequest = !event.responseHeaders
+            const isRequest = !event.responseHeaders && !event.responseErrorReason
 
-            event.responseStatusCode ||= 200
+            event.responseStatusCode ||= event.responseErrorReason ? 0 : 200
             event.responseHeaders ||= []
 
             const responseHeaders = event.responseHeaders.reduce((headers, { name, value }) => {
@@ -220,10 +224,27 @@ export default class DevtoolsInterception extends Interception {
     /**
      * Does everything that `mock.clear()` does, and also
      * removes any mocked return values or implementations.
+     * Restored mock does not emit events and could not mock responses
      */
-    restore () {
+    async restore (sessionMocks = SESSION_MOCKS, cdpSessions = CDP_SESSIONS) {
         this.clear()
         this.respondOverwrites = []
+        this.restored = true
+        const handle = await this.browser.getWindowHandle()
+
+        log.trace(`Restoring mock for ${handle}`)
+        sessionMocks[handle].delete(this)
+
+        if (sessionMocks[handle].size) {
+            return
+        }
+
+        log.trace(`Disabling fetch domain for ${handle}`)
+        return cdpSessions[handle].send('Fetch.disable')
+            .then(() => {
+                delete sessionMocks[handle]
+                delete cdpSessions[handle]
+            }).catch(/* istanbul ignore next */logFetchError)
     }
 
     /**
@@ -232,6 +253,7 @@ export default class DevtoolsInterception extends Interception {
      * @param {*} params      additional respond parameters to overwrite
      */
     respond (overwrite: MockOverwrite, params: MockResponseParams = {}) {
+        this.ensureNotRestored()
         this.respondOverwrites.push({ overwrite, params, sticky: true })
     }
 
@@ -241,6 +263,7 @@ export default class DevtoolsInterception extends Interception {
      * @param {*} params      additional respond parameters to overwrite
      */
     respondOnce (overwrite: MockOverwrite, params: MockResponseParams = {}) {
+        this.ensureNotRestored()
         this.respondOverwrites.push({ overwrite, params })
     }
 
@@ -249,6 +272,7 @@ export default class DevtoolsInterception extends Interception {
      * @param {string} errorCode  error code of the response
      */
     abort (errorReason: Protocol.Network.ErrorReason, sticky: boolean = true) {
+        this.ensureNotRestored()
         if (typeof errorReason !== 'string' || !ERROR_REASON.includes(errorReason)) {
             throw new Error(`Invalid value for errorReason, allowed are: ${ERROR_REASON.join(', ')}`)
         }
@@ -262,6 +286,12 @@ export default class DevtoolsInterception extends Interception {
     abortOnce (errorReason: Protocol.Network.ErrorReason) {
         this.abort(errorReason, false)
     }
+
+    private ensureNotRestored() {
+        if (this.restored) {
+            throw new Error('This can\'t be done on restored mock')
+        }
+    }
 }
 
 const filterMethod = (method: string, expected?: ExpectParameter<string>) => {

packages/webdriverio/src/utils/interception/index.ts

@@ -11,7 +11,7 @@ import type { Protocol } from 'devtools-protocol'
 export default abstract class Interception extends EventEmitter {
     abstract calls: Matches[] | Promise<Matches[]>
     abstract clear (): void
-    abstract restore (): void
+    abstract restore (): Promise<void>
     abstract respond (overwrite: MockOverwrite, params: MockResponseParams): void
     abstract respondOnce (overwrite: MockOverwrite, params: MockResponseParams): void
     abstract abort (errorReason: Protocol.Network.ErrorReason, sticky: boolean): void

packages/webdriverio/src/utils/interception/webdriver.ts

@@ -47,8 +47,8 @@ export default class WebDriverInterception extends Interception {
      * Does everything that `mock.clear()` does, and also
      * removes any mocked return values or implementations.
      */
-    restore () {
-        return this.browser.call(
+    async restore () {
+        await this.browser.call(
             async () => this.browser.clearMockCalls(this.mockId as string, true))
     }
 

packages/webdriverio/tests/utils/interception/devtools.test.ts

@@ -232,6 +232,18 @@ test('decodes base64 responses', async () => {
     expect(mock.calls[1].body).toBe('{"foo":"bar"}')
 })
 
+test('response error', async () => {
+    const mock = new NetworkInterception('**/foobar/**', undefined, browserMock)
+    cdpClient.send.mockReturnValueOnce(Promise.resolve({}))
+    cdpClient.send.mockReturnValueOnce(Promise.resolve({}))
+    await fetchListener(mock, {
+        request: { url: 'http://test.com/foobar/test.html' },
+        responseErrorReason: 'NameNotResolved'
+    })
+
+    expect(mock.calls[0].statusCode).toBe(0)
+})
+
 test('undefined response', async () => {
     const mock = new NetworkInterception('**/foobar/**', undefined, browserMock)
     cdpClient.send.mockReturnValueOnce(Promise.resolve({}))
@@ -645,7 +657,12 @@ test('allows to clear mocks', async () => {
 })
 
 test('allows to restore mocks', async () => {
+    const browserMock = {
+        getWindowHandle: vi.fn().mockResolvedValue('window-handle')
+    } as any as Browser
     const mock = new NetworkInterception('**/foobar/**', undefined, browserMock)
+    const sessionMocks = { 'window-handle': new Set([mock]) }
+    const cdpSessions = { 'window-handle': cdpClient }
     mock.respondOnce({ foo: 'bar' })
     mock.respond({ bar: 'foo' })
 
@@ -661,6 +678,49 @@ test('allows to restore mocks', async () => {
     })
     expect(mock.respondOverwrites.length).toBe(1)
 
-    mock.restore()
+    mock.restore(sessionMocks, cdpSessions)
     expect(mock.respondOverwrites.length).toBe(0)
 })
+
+test('removes mock after restore', async () => {
+    const browserMock = {
+        getWindowHandle: vi.fn().mockResolvedValue('window-handle')
+    } as any as Browser
+    const mock = new NetworkInterception('**/foobar/**', undefined, browserMock)
+    const sessionMocks = { 'window-handle': new Set([mock]) }
+    const cdpSessions = { 'window-handle': cdpClient }
+
+    await mock.restore(sessionMocks, cdpSessions)
+
+    expect(sessionMocks).toEqual({})
+})
+
+test('disables fetch domain after restore, if there are no other mocks', async () => {
+    const browserMock = {
+        getWindowHandle: vi.fn().mockResolvedValue('window-handle')
+    } as any as Browser
+    const mock = new NetworkInterception('**/foobar/**', undefined, browserMock)
+    const sessionMocks = { 'window-handle': new Set([mock]) }
+    const cdpSessions = { 'window-handle': cdpClient }
+
+    await mock.restore(sessionMocks, cdpSessions)
+
+    expect(cdpClient.send).toBeCalledWith('Fetch.disable')
+    expect(cdpSessions).toEqual({})
+})
+
+test('does not disable fetch domain after restore, if there are other mocks', async () => {
+    const browserMock = {
+        getWindowHandle: vi.fn().mockResolvedValue('window-handle')
+    } as any as Browser
+    const firstMock = new NetworkInterception('**/foobar/**', undefined, browserMock)
+    const secondMock = new NetworkInterception('**/foobar/**', undefined, browserMock)
+    const sessionMocks = { 'window-handle': new Set([firstMock, secondMock]) }
+    const cdpSessions = { 'window-handle': cdpClient }
+
+    await firstMock.restore(sessionMocks, cdpSessions)
+
+    expect(cdpClient.send).not.toBeCalledWith('Fetch.disable')
+    expect(sessionMocks).toEqual({ 'window-handle': new Set([secondMock]) })
+    expect(cdpSessions).toEqual({ 'window-handle': cdpClient })
+})

packages/webdriverio/tests/utils/interception/webdriver.test.ts

@@ -47,7 +47,7 @@ test('clear', async () => {
 test('restore', async () => {
     const mock = new NetworkInterception('**/foobar/**', {}, browserMock)
     await mock.init()
-    expect(await mock.restore()).toEqual({})
+    expect(await mock.restore()).toBeUndefined()
     expect(browserMock.clearMockCalls).toBeCalledWith(123, true)
 })