fix(types): Fix overwriteCommand inconsistent typing (#14549)

dprevost-LMI · web-flow · commit ce33265ece67 · 2025-06-09T17:34:58.000-07:00
* Fix incorrect inference when overriding commands for element and browser

- Fix inference having the `this` which was in excess
- Fix inference wrapping in a `(...arg) =&gt;` which was in excess
- Fix browser command overwrite where the return promise type is wrong because undefined verbose return type is inferred to `Promise&lt;unknown&gt;`

* Fix the debug return type

* Fix typings

* Fix UTs

* Have more verbose and type safe examples

* fix the type reference of the `this` argument + add UTs

* Even better example

* More type fixes + more robust UTs

* Simplification of code in index.ts

* Add this infer type for browser overwritten command + final review

* Remove proposition to revisite later + fix typing issue
diff --git a/packages/webdriverio/src/commands/browser/custom$.ts b/packages/webdriverio/src/commands/browser/custom$.ts
@@ -34,7 +34,7 @@ export async function custom$ (
     this: WebdriverIO.Browser,
     strategyName: string,
     ...strategyArguments: unknown[]
-) {
+) : Promise<WebdriverIO.Element> {
     const strategy = this.strategies.get(strategyName) as CustomStrategyFunction
 
     if (!strategy) {
diff --git a/packages/webdriverio/src/commands/browser/debug.ts b/packages/webdriverio/src/commands/browser/debug.ts
@@ -34,7 +34,7 @@ import WDIORepl from '@wdio/repl'
 export function debug(
     this: WebdriverIO.Browser,
     commandTimeout = 5000
-) {
+): Promise<void | unknown> {
     const repl = new WDIORepl()
     const { introMessage } = WDIORepl
 
diff --git a/packages/webdriverio/src/commands/browser/getWindowSize.ts b/packages/webdriverio/src/commands/browser/getWindowSize.ts
@@ -23,9 +23,9 @@ interface BrowserSize {
  * @type window
  *
  */
-export async function getWindowSize(this: WebdriverIO.Browser) {
+export async function getWindowSize(this: WebdriverIO.Browser): Promise<BrowserSize> {
     const browser = getBrowserObject(this)
 
     const { width, height } = await browser.getWindowRect() as BrowserSize
-    return { width, height } as BrowserSize
+    return { width, height } satisfies BrowserSize
 }
diff --git a/packages/webdriverio/src/commands/browser/keys.ts b/packages/webdriverio/src/commands/browser/keys.ts
@@ -26,7 +26,7 @@ import { checkUnicode } from '../../utils/index.js'
 export async function keys (
     this: WebdriverIO.Browser,
     value: string | string[]
-) {
+): Promise<void> {
     let keySequence: string[] = []
 
     /**
diff --git a/packages/webdriverio/src/commands/browser/mockClearAll.ts b/packages/webdriverio/src/commands/browser/mockClearAll.ts
@@ -30,7 +30,7 @@ const log = logger('webdriverio:mockClearAll')
  *
  * @alias browser.mockClearAll
  */
-export async function mockClearAll () {
+export async function mockClearAll (): Promise<void> {
     for (const [handle, mocks] of Object.entries(SESSION_MOCKS)) {
         log.trace(`Clearing mocks for ${handle}`)
         for (const mock of mocks) {
diff --git a/packages/webdriverio/src/commands/browser/mockRestoreAll.ts b/packages/webdriverio/src/commands/browser/mockRestoreAll.ts
@@ -28,7 +28,7 @@ const log = logger('webdriverio:mockRestoreAll')
  *
  * @alias browser.mockRestoreAll
  */
-export async function mockRestoreAll () {
+export async function mockRestoreAll (): Promise<void> {
     for (const [handle, mocks] of Object.entries(SESSION_MOCKS)) {
         log.trace(`Clearing mocks for ${handle}`)
         for (const mock of mocks) {
diff --git a/packages/webdriverio/src/commands/browser/pause.ts b/packages/webdriverio/src/commands/browser/pause.ts
@@ -22,6 +22,6 @@
 export function pause (
     this: WebdriverIO.Browser,
     milliseconds = 1000
-) {
+): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, milliseconds))
 }
diff --git a/packages/webdriverio/src/commands/browser/react$$.ts b/packages/webdriverio/src/commands/browser/react$$.ts
@@ -43,7 +43,7 @@ export async function react$$ (
     this: WebdriverIO.Browser,
     selector: string,
     { props = {}, state = {} }: ReactSelectorOptions = {}
-) {
+): Promise<WebdriverIO.ElementArray> {
     await this.executeScript(resqScript, [])
     await this.execute(waitToLoadReact)
     const res = await this.execute(
diff --git a/packages/webdriverio/src/commands/browser/react$.ts b/packages/webdriverio/src/commands/browser/react$.ts
@@ -52,7 +52,7 @@ export async function react$ (
     this: WebdriverIO.Browser,
     selector: string,
     { props = {}, state = {} }: ReactSelectorOptions = {}
-) {
+): Promise<WebdriverIO.Element> {
     await this.executeScript(resqScript.toString(), [])
     await this.execute(waitToLoadReact)
     const res = await this.execute(
diff --git a/packages/webdriverio/src/commands/browser/reloadSession.ts b/packages/webdriverio/src/commands/browser/reloadSession.ts
@@ -55,7 +55,7 @@ const log = logger('webdriverio')
  * @type utility
  *
  */
-export async function reloadSession (this: WebdriverIO.Browser, newCapabilities?: WebdriverIO.Capabilities) {
+export async function reloadSession (this: WebdriverIO.Browser, newCapabilities?: WebdriverIO.Capabilities): Promise<string> {
     const oldSessionId = (this as WebdriverIO.Browser).sessionId
 
     /**
@@ -91,5 +91,5 @@ export async function reloadSession (this: WebdriverIO.Browser, newCapabilities?
         await Promise.all(options.onReload.map((hook) => hook(oldSessionId, (this as WebdriverIO.Browser).sessionId)))
     }
 
-    return this.sessionId as string
+    return this.sessionId
 }
diff --git a/packages/webdriverio/src/commands/browser/restore.ts b/packages/webdriverio/src/commands/browser/restore.ts
@@ -34,7 +34,7 @@ import { restoreFunctions } from '../../constants.js'
 export async function restore (
     this: WebdriverIO.Browser,
     scopes?: SupportedScopes | SupportedScopes[]
-) {
+): Promise<void> {
     const scopeArray = !scopes || Array.isArray(scopes) ? scopes : [scopes]
     const instanceRestoreFunctions = restoreFunctions.get(this)
     if (!instanceRestoreFunctions) {
diff --git a/packages/webdriverio/src/commands/browser/savePDF.ts b/packages/webdriverio/src/commands/browser/savePDF.ts
@@ -47,7 +47,7 @@ export async function savePDF (
     this: WebdriverIO.Browser,
     filepath: string,
     options?: PDFPrintOptions
-) {
+): Promise<Buffer<ArrayBuffer>> {
     /**
      * run command implementation based on given environment
      */
diff --git a/packages/webdriverio/src/commands/browser/saveRecordingScreen.ts b/packages/webdriverio/src/commands/browser/saveRecordingScreen.ts
@@ -28,7 +28,7 @@ import { environment } from '../../environment.js'
 export async function saveRecordingScreen (
     this: WebdriverIO.Browser,
     filepath: string
-) {
+): Promise<Buffer<ArrayBuffer>> {
     /**
      * run command implementation based on given environment
      */
diff --git a/packages/webdriverio/src/commands/browser/saveScreenshot.ts b/packages/webdriverio/src/commands/browser/saveScreenshot.ts
@@ -52,7 +52,7 @@ export async function saveScreenshot (
     this: WebdriverIO.Browser,
     filepath: string,
     options?: SaveScreenshotOptions
-) {
+): Promise<Buffer<ArrayBuffer>> {
     /**
      * run command implementation based on given environment
      */
diff --git a/packages/webdriverio/src/commands/browser/scroll.ts b/packages/webdriverio/src/commands/browser/scroll.ts
@@ -29,9 +29,10 @@ export function scroll (
     this: WebdriverIO.Browser,
     x = 0,
     y = 0
-) {
+): Promise<void> {
     if (!x && !y) {
-        return log.warn('"scroll" command was called with no parameters, skipping execution')
+        log.warn('"scroll" command was called with no parameters, skipping execution')
+        return Promise.resolve()
     }
 
     // Appium does not support the "wheel" action
diff --git a/packages/webdriverio/src/commands/browser/setCookies.ts b/packages/webdriverio/src/commands/browser/setCookies.ts
@@ -57,7 +57,7 @@ import type { Cookie } from '@wdio/protocols'
 export async function setCookies(
     this: WebdriverIO.Browser,
     cookieObjs: Cookie | Cookie[]
-) {
+): Promise<void> {
     const cookieObjsList = !Array.isArray(cookieObjs) ? [cookieObjs] : cookieObjs
 
     if (cookieObjsList.some(obj => (typeof obj !== 'object'))) {
diff --git a/packages/webdriverio/src/commands/browser/setViewport.ts b/packages/webdriverio/src/commands/browser/setViewport.ts
@@ -36,7 +36,7 @@ export interface SetViewportOptions {
 export async function setViewport(
     this: WebdriverIO.Browser,
     options: SetViewportOptions
-) {
+): Promise<void> {
     /**
      * type check
      */
diff --git a/packages/webdriverio/src/commands/browser/setWindowSize.ts b/packages/webdriverio/src/commands/browser/setWindowSize.ts
@@ -27,7 +27,7 @@ export async function setWindowSize(
     this: WebdriverIO.Browser,
     width: number,
     height: number
-) {
+): Promise<void> {
     /**
      * type check
      */
diff --git a/packages/webdriverio/src/commands/browser/switchFrame.ts b/packages/webdriverio/src/commands/browser/switchFrame.ts
@@ -68,7 +68,7 @@ const log = logger('webdriverio:switchFrame')
 export async function switchFrame (
     this: WebdriverIO.Browser,
     context: WebdriverIO.Element | ChainablePromiseElement | string | null | ((tree: FlatContextTree) => boolean | Promise<boolean>)
-) {
+): Promise<string | void> {
     function isPossiblyUnresolvedElement(input: typeof context): input is WebdriverIO.Element | ChainablePromiseElement {
         return Boolean(input) && typeof input === 'object' && typeof (input as WebdriverIO.Element).getElement === 'function'
     }
diff --git a/packages/webdriverio/src/commands/browser/switchWindow.ts b/packages/webdriverio/src/commands/browser/switchWindow.ts
@@ -37,7 +37,7 @@ import { getContextManager } from '../../session/context.js'
 export async function switchWindow (
     this: WebdriverIO.Browser,
     matcher: string | RegExp
-) {
+): Promise<string> {
     /**
      * parameter check
      */
diff --git a/packages/webdriverio/src/commands/browser/throttle.ts b/packages/webdriverio/src/commands/browser/throttle.ts
@@ -10,7 +10,7 @@ const log = logger('webdriverio:throttle')
 export async function throttle (
     this: WebdriverIO.Browser,
     params: ThrottleOptions
-) {
+): Promise<void> {
     log.warn('Command "throttle" is deprecated and will be removed with the next major version release! Use `throttleNetwork` instead.')
     const browser = getBrowserObject(this)
     await browser.throttleNetwork(params)
diff --git a/packages/webdriverio/src/commands/browser/throttleCPU.ts b/packages/webdriverio/src/commands/browser/throttleCPU.ts
@@ -26,7 +26,7 @@
 export async function throttleCPU (
     this: WebdriverIO.Browser,
     factor: number
-) {
+): Promise<void> {
     if (typeof factor !== 'number') {
         throw new Error('Invalid factor for "throttleCPU". Expected it to be a number (int)')
     }
diff --git a/packages/webdriverio/src/commands/browser/throttleNetwork.ts b/packages/webdriverio/src/commands/browser/throttleNetwork.ts
@@ -114,7 +114,7 @@ const NETWORK_PRESET_TYPES = Object.keys(NETWORK_PRESETS)
 export async function throttleNetwork (
     this: WebdriverIO.Browser,
     params: ThrottleOptions
-) {
+): Promise<void> {
     if (
         /**
          * check string parameter
@@ -134,7 +134,7 @@ export async function throttleNetwork (
     if (this.isSauce) {
         const browser = getBrowserObject(this)
         await browser.sauceThrottleNetwork(params)
-        return null
+        return
     }
 
     const failedConnectionMessage = 'No Puppeteer connection could be established which is required to use this command'
@@ -161,5 +161,5 @@ export async function throttleNetwork (
             : params
     )
 
-    return null
+    return
 }
diff --git a/packages/webdriverio/src/commands/browser/touchAction.ts b/packages/webdriverio/src/commands/browser/touchAction.ts
@@ -75,6 +75,6 @@ import type { TouchActions } from '../../types.js'
 export function touchAction (
     this: WebdriverIO.Browser,
     actions: TouchActions
-) {
+): Promise<void> {
     return touchActionCommand.call(this, actions)
 }
diff --git a/packages/webdriverio/src/index.ts b/packages/webdriverio/src/index.ts
@@ -68,15 +68,9 @@ export const remote = async function(
      * we need to overwrite the original addCommand and overwriteCommand
      */
     if ((params as Options.Testrunner).framework && !isStub(params.automationProtocol)) {
-        const origAddCommand = instance.addCommand.bind(instance) as typeof instance.addCommand
-        instance.addCommand = (name: string, fn: (...args: any[]) => any, attachToElement) => (
-            origAddCommand(name, fn, attachToElement)
-        )
-
-        const origOverwriteCommand = instance.overwriteCommand.bind(instance) as typeof instance.overwriteCommand
-        instance.overwriteCommand = (name: string, fn: (...args: any[]) => any, attachToElement) => (
-            origOverwriteCommand<keyof typeof elementCommands, any, any>(name, fn, attachToElement)
-        )
+        instance.addCommand = instance.addCommand.bind(instance)
+
+        instance.overwriteCommand = instance.overwriteCommand.bind(instance)
     }
 
     instance.addLocatorStrategy = addLocatorStrategyHandler(instance)
@@ -178,9 +172,9 @@ export const multiremote = async function (
      */
     if (!isStub(automationProtocol)) {
         const origAddCommand = driver.addCommand.bind(driver)
-        driver.addCommand = (name: string, fn: (...args: any[]) => any, attachToElement) => {
-            driver.instances.forEach(instance =>
-                driver.getInstance(instance).addCommand(name, fn, attachToElement)
+        driver.addCommand = (name: string, fn, attachToElement) => {
+            driver.instances.forEach(instanceName =>
+                driver.getInstance(instanceName).addCommand(name, fn, attachToElement)
             )
 
             return origAddCommand(
@@ -193,7 +187,7 @@ export const multiremote = async function (
         }
 
         const origOverwriteCommand = driver.overwriteCommand.bind(driver) as typeof driver.overwriteCommand
-        driver.overwriteCommand = (name: string, fn: (...args: any[]) => any, attachToElement) => {
+        driver.overwriteCommand = (name, fn, attachToElement) => {
             return origOverwriteCommand<keyof typeof elementCommands, any, any>(
                 name,
                 fn,
diff --git a/packages/webdriverio/src/types.ts b/packages/webdriverio/src/types.ts
@@ -189,7 +189,7 @@ type AddCommandFnScoped<
     InstanceType = WebdriverIO.Browser,
     IsElement extends boolean = false
 > = (
-    this: IsElement extends true ? Element : InstanceType,
+    this: IsElement extends true ? WebdriverIO.Element : InstanceType,
     ...args: any[]
 ) => any
 
@@ -201,7 +201,7 @@ type OverwriteCommandFnScoped<
     IsElement extends boolean = false
 > = (
     this: IsElement extends true ? WebdriverIO.Element : WebdriverIO.Browser,
-    origCommand: (...args: any[]) => IsElement extends true ? $ElementCommands[ElementKey] : $BrowserCommands[BrowserKey],
+    originalCommand: IsElement extends true ? OmitThisParameter<$ElementCommands[ElementKey]> : OmitThisParameter<$BrowserCommands[BrowserKey]>,
     ...args: any[]
 ) => Promise<any>
 
@@ -210,32 +210,36 @@ type OverwriteCommandFn<
     BrowserKey extends keyof $BrowserCommands,
     IsElement extends boolean = false
 > = (
-    origCommand: (...args: any[]) => IsElement extends true ? $ElementCommands[ElementKey] : $BrowserCommands[BrowserKey],
+    this: IsElement extends true ? WebdriverIO.Element : WebdriverIO.Browser,
+    originalCommand: IsElement extends true ? OmitThisParameter<$ElementCommands[ElementKey]> : OmitThisParameter<$BrowserCommands[BrowserKey]>,
     ...args: any[]
 ) => Promise<any>
 
 export type CustomLocatorReturnValue = HTMLElement | HTMLElement[] | NodeListOf<HTMLElement>
+
+type Instances = WebdriverIO.Browser | WebdriverIO.MultiRemoteBrowser
+
 export interface CustomInstanceCommands<T> {
     /**
      * add command to `browser` or `element` scope
      */
-    addCommand<IsElement extends boolean = false>(
+    addCommand<IsElement extends boolean = false, Instance extends Instances = WebdriverIO.Browser>(
         name: string,
-        func: AddCommandFn | AddCommandFnScoped<T, IsElement>,
+        func: IsElement extends true ? AddCommandFnScoped<T | Instance, IsElement> : AddCommandFn,
         attachToElement?: IsElement,
         proto?: Record<string, any>,
-        instances?: Record<string, WebdriverIO.Browser | WebdriverIO.MultiRemoteBrowser>
+        instances?: Record<string, Instances>
     ): void;
 
     /**
      * overwrite `browser` or `element` command
      */
     overwriteCommand<ElementKey extends keyof $ElementCommands, BrowserKey extends keyof $BrowserCommands, IsElement extends boolean = false>(
         name: IsElement extends true ? ElementKey : BrowserKey,
-        func: OverwriteCommandFn<ElementKey, BrowserKey, IsElement> | OverwriteCommandFnScoped<ElementKey, BrowserKey, IsElement>,
+        func: IsElement extends true ? OverwriteCommandFnScoped<ElementKey, BrowserKey, IsElement> : OverwriteCommandFn<ElementKey, BrowserKey, IsElement>,
         attachToElement?: IsElement,
         proto?: Record<string, any>,
-        instances?: Record<string, WebdriverIO.Browser | WebdriverIO.MultiRemoteBrowser>
+        instances?: Record<string, Instances>
     ): void;
 
     /**
diff --git a/packages/webdriverio/tests/addCommand.test.ts b/packages/webdriverio/tests/addCommand.test.ts
@@ -59,6 +59,26 @@ const customCommand = async () => {
 
 describe('addCommand', () => {
     describe('remote', () => {
+
+        test('should resolve the this parameter by inference', async () => {
+            const browser = await remote(remoteConfig)
+            browser.addCommand(
+                'press',
+                async function (this /* Expect to be infer to WebDriverIO.Element by default */) {
+                    await this.click()
+                    return
+                },
+                true,
+            )
+
+            const element = await browser.$('.someRandomElement')
+            vi.spyOn(element, 'click')
+
+            // @ts-expect-error undefined custom command
+            expect(await element.press()).toBeUndefined()
+            expect(element.click).toBeCalledTimes(1)
+        })
+
         test('should be able to handle async', async () => {
             const browser = await remote(remoteConfig)
 
diff --git a/packages/webdriverio/tests/overwriteCommand.test.ts b/packages/webdriverio/tests/overwriteCommand.test.ts
diff --git a/website/docs/CustomCommands.md b/website/docs/CustomCommands.md

Original file line number	Diff line number	Diff line change
`@@ -23,9 +23,9 @@ interface BrowserSize {`
`23`	`23`	`* @type window`
`24`	`24`	`*`
`25`	`25`	`*/`
`26`		`-export async function getWindowSize(this: WebdriverIO.Browser) {`
	`26`	`+export async function getWindowSize(this: WebdriverIO.Browser): Promise<BrowserSize> {`
`27`	`27`	`const browser = getBrowserObject(this)`
`28`	`28`
`29`	`29`	`const { width, height } = await browser.getWindowRect() as BrowserSize`
`30`		`- return { width, height } as BrowserSize`
	`30`	`+ return { width, height } satisfies BrowserSize`
`31`	`31`	`}`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ const log = logger('webdriverio:mockClearAll')`
`30`	`30`	`*`
`31`	`31`	`* @alias browser.mockClearAll`
`32`	`32`	`*/`
`33`		`-export async function mockClearAll () {`
	`33`	`+export async function mockClearAll (): Promise<void> {`
`34`	`34`	`for (const [handle, mocks] of Object.entries(SESSION_MOCKS)) {`
`35`	`35`	log.trace(`Clearing mocks for ${handle}`)
`36`	`36`	`for (const mock of mocks) {`
Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@ const log = logger('webdriverio:mockRestoreAll')`
`28`	`28`	`*`
`29`	`29`	`* @alias browser.mockRestoreAll`
`30`	`30`	`*/`
`31`		`-export async function mockRestoreAll () {`
	`31`	`+export async function mockRestoreAll (): Promise<void> {`
`32`	`32`	`for (const [handle, mocks] of Object.entries(SESSION_MOCKS)) {`
`33`	`33`	log.trace(`Clearing mocks for ${handle}`)
`34`	`34`	`for (const mock of mocks) {`
Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,6 @@`
`22`	`22`	`export function pause (`
`23`	`23`	`this: WebdriverIO.Browser,`
`24`	`24`	`milliseconds = 1000`
`25`		`-) {`
	`25`	`+): Promise<void> {`
`26`	`26`	`return new Promise((resolve) => setTimeout(resolve, milliseconds))`
`27`	`27`	`}`
Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ const log = logger('webdriverio')`
`55`	`55`	`* @type utility`
`56`	`56`	`*`
`57`	`57`	`*/`
`58`		`-export async function reloadSession (this: WebdriverIO.Browser, newCapabilities?: WebdriverIO.Capabilities) {`
	`58`	`+export async function reloadSession (this: WebdriverIO.Browser, newCapabilities?: WebdriverIO.Capabilities): Promise<string> {`
`59`	`59`	`const oldSessionId = (this as WebdriverIO.Browser).sessionId`
`60`	`60`
`61`	`61`	`/**`
`@@ -91,5 +91,5 @@ export async function reloadSession (this: WebdriverIO.Browser, newCapabilities?`
`91`	`91`	`await Promise.all(options.onReload.map((hook) => hook(oldSessionId, (this as WebdriverIO.Browser).sessionId)))`
`92`	`92`	`}`
`93`	`93`
`94`		`- return this.sessionId as string`
	`94`	`+ return this.sessionId`
`95`	`95`	`}`