fix(wdio): wait for the driver to be properly closed (#8305)

sheremet-va · web-flow · commit c16abe71eaae · 2025-07-22T15:41:50.000+02:00
diff --git a/docs/guide/browser/commands.md b/docs/guide/browser/commands.md
@@ -162,7 +162,7 @@ Vitest exposes some `webdriverio` specific properties on the context object.
 
 - `browser` is the `WebdriverIO.Browser` API.
 
-Vitest automatically switches the `webdriver` context to the test iframe by calling `browser.switchToFrame` before the command is called, so `$` and `$$` methods refer to the elements inside the iframe, not in the orchestrator, but non-webdriver APIs will still refer to the parent frame context.
+Vitest automatically switches the `webdriver` context to the test iframe by calling `browser.switchFrame` before the command is called, so `$` and `$$` methods refer to the elements inside the iframe, not in the orchestrator, but non-webdriver APIs will still refer to the parent frame context.
 
 ::: tip
 If you are using TypeScript, don't forget to reference `@vitest/browser/providers/webdriverio` in your [setup file](/config/#setupfile) or a [config file](/config/) to get autocompletion:
diff --git a/packages/browser/src/node/plugin.ts b/packages/browser/src/node/plugin.ts
@@ -385,7 +385,7 @@ export default (parentServer: ParentBrowserProject, base = '/'): Plugin[] => {
           viteConfig.esbuild.legalComments = 'inline'
         }
 
-        const defaultPort = parentServer.vitest._browserLastPort++
+        const defaultPort = parentServer.vitest.state._data.browserLastPort++
 
         const api = resolveApiServerConfig(
           viteConfig.test?.browser || {},
diff --git a/packages/browser/src/node/pool.ts b/packages/browser/src/node/pool.ts
@@ -101,16 +101,17 @@ export function createBrowserPool(vitest: Vitest): ProcessPool {
     const parallelPools: (() => Promise<void>)[] = []
     const nonParallelPools: (() => Promise<void>)[] = []
 
-    for (const result of initialisedPools) {
-      if (!result) {
+    for (const pool of initialisedPools) {
+      if (!pool) {
+        // this means it was cancelled
         return
       }
 
-      if (result.provider.mocker && result.provider.supportsParallelism) {
-        parallelPools.push(result.runTests)
+      if (pool.provider.mocker && pool.provider.supportsParallelism) {
+        parallelPools.push(pool.runTests)
       }
       else {
-        nonParallelPools.push(result.runTests)
+        nonParallelPools.push(pool.runTests)
       }
     }
 
diff --git a/packages/browser/src/node/providers/webdriver.ts b/packages/browser/src/node/providers/webdriver.ts
@@ -36,11 +36,17 @@ export class WebdriverBrowserProvider implements BrowserProvider {
   }
 
   async initialize(
-    ctx: TestProject,
+    project: TestProject,
     { browser, options }: WebdriverProviderOptions,
   ): Promise<void> {
+    // increase shutdown timeout because WDIO takes some extra time to kill the driver
+    if (!project.vitest.state._data.timeoutIncreased) {
+      project.vitest.state._data.timeoutIncreased = true
+      project.vitest.config.teardownTimeout += 10_000
+    }
+
     this.closing = false
-    this.project = ctx
+    this.project = project
     this.browserName = browser
     this.options = options as Capabilities.WebdriverIOConfig
   }
@@ -50,17 +56,17 @@ export class WebdriverBrowserProvider implements BrowserProvider {
   }
 
   async switchToTestFrame(): Promise<void> {
-    const page = this.browser!
+    const browser = this.browser!
     // support wdio@9
-    if (page.switchFrame) {
-      await page.switchFrame(page.$('iframe[data-vitest]'))
+    if (browser.switchFrame) {
+      await browser.switchFrame(browser.$('iframe[data-vitest]'))
     }
     else {
-      const iframe = await page.findElement(
+      const iframe = await browser.findElement(
         'css selector',
         'iframe[data-vitest]',
       )
-      await page.switchToFrame(iframe)
+      await browser.switchToFrame(iframe)
     }
     this.iframeSwitched = true
   }
@@ -190,11 +196,15 @@ export class WebdriverBrowserProvider implements BrowserProvider {
   async close(): Promise<void> {
     debug?.('[%s] closing provider', this.browserName)
     this.closing = true
-    await Promise.all([
-      this.browser?.sessionId ? this.browser?.deleteSession?.() : null,
-    ])
-    // TODO: right now process can only exit with timeout, if we use browser
-    // needs investigating
-    process.exit()
+    const browser = this.browser
+    const sessionId = browser?.sessionId
+    if (!browser || !sessionId) {
+      return
+    }
+
+    // https://github.com/webdriverio/webdriverio/blob/ab1a2e82b13a9c7d0e275ae87e7357e1b047d8d3/packages/wdio-runner/src/index.ts#L486
+    await browser.deleteSession()
+    browser.sessionId = undefined as unknown as string
+    this.browser = null
   }
 }
diff --git a/packages/browser/src/node/rpc.ts b/packages/browser/src/node/rpc.ts
@@ -58,15 +58,17 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke
       )
     }
 
-    if (!vitest._browserSessions.sessionIds.has(sessionId)) {
-      const ids = [...vitest._browserSessions.sessionIds].join(', ')
+    const sessions = vitest._browserSessions
+
+    if (!sessions.sessionIds.has(sessionId)) {
+      const ids = [...sessions.sessionIds].join(', ')
       return error(
         new Error(`[vitest] Unknown session id "${sessionId}". Expected one of ${ids}.`),
       )
     }
 
     if (type === 'orchestrator') {
-      const session = vitest._browserSessions.getSession(sessionId)
+      const session = sessions.getSession(sessionId)
       // it's possible the session was already resolved by the preview provider
       session?.connected()
     }
@@ -94,7 +96,7 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke
         clients.delete(rpcId)
         globalServer.removeCDPHandler(rpcId)
         if (type === 'orchestrator') {
-          vitest._browserSessions.destroySession(sessionId)
+          sessions.destroySession(sessionId)
         }
         // this will reject any hanging methods if there are any
         rpc.$close(
diff --git a/packages/vitest/src/node/core.ts b/packages/vitest/src/node/core.ts
@@ -19,7 +19,6 @@ import { ViteNodeRunner } from 'vite-node/client'
 import { ViteNodeServer } from 'vite-node/server'
 import { version } from '../../package.json' with { type: 'json' }
 import { WebSocketReporter } from '../api/setup'
-import { defaultBrowserPort } from '../constants'
 import { distDir } from '../paths'
 import { wildcardPatternToRegExp } from '../utils/base'
 import { convertTasksToEvents } from '../utils/tasks'
@@ -89,7 +88,6 @@ export class Vitest {
   /** @internal */ closingPromise?: Promise<void>
   /** @internal */ isCancelling = false
   /** @internal */ coreWorkspaceProject: TestProject | undefined
-  /** @internal */ _browserLastPort = defaultBrowserPort
   /** @internal */ _browserSessions = new BrowserSessions()
   /** @internal */ _cliOptions: CliOptions = {}
   /** @internal */ reporters: Reporter[] = []
@@ -196,7 +194,6 @@ export class Vitest {
     this.watcher.unregisterWatcher()
     clearTimeout(this._rerunTimer)
     this.restartsCount += 1
-    this._browserLastPort = defaultBrowserPort
     this.pool?.close?.()
     this.pool = undefined
     this.closingPromise = undefined
diff --git a/packages/vitest/src/node/state.ts b/packages/vitest/src/node/state.ts
@@ -4,6 +4,7 @@ import type { TestProject } from './project'
 import type { MergedBlobs } from './reporters/blob'
 import type { OnUnhandledErrorCallback } from './types/config'
 import { createFileTask } from '@vitest/runner/utils'
+import { defaultBrowserPort } from '../constants'
 import { TestCase, TestModule, TestSuite } from './reporters/reported-tasks'
 
 function isAggregateError(err: unknown): err is AggregateError {
@@ -26,6 +27,12 @@ export class StateManager {
 
   onUnhandledError?: OnUnhandledErrorCallback
 
+  /** @internal */
+  _data = {
+    browserLastPort: defaultBrowserPort,
+    timeoutIncreased: false,
+  }
+
   constructor(
     options: {
       onUnhandledError?: OnUnhandledErrorCallback

Original file line number	Diff line number	Diff line change
`@@ -385,7 +385,7 @@ export default (parentServer: ParentBrowserProject, base = '/'): Plugin[] => {`
`385`	`385`	`viteConfig.esbuild.legalComments = 'inline'`
`386`	`386`	`}`
`387`	`387`
`388`		`- const defaultPort = parentServer.vitest._browserLastPort++`
	`388`	`+ const defaultPort = parentServer.vitest.state._data.browserLastPort++`
`389`	`389`
`390`	`390`	`const api = resolveApiServerConfig(`
`391`	`391`	`viteConfig.test?.browser \|\| {},`
Original file line number	Diff line number	Diff line change
`@@ -101,16 +101,17 @@ export function createBrowserPool(vitest: Vitest): ProcessPool {`
`101`	`101`	`const parallelPools: (() => Promise<void>)[] = []`
`102`	`102`	`const nonParallelPools: (() => Promise<void>)[] = []`
`103`	`103`
`104`		`- for (const result of initialisedPools) {`
`105`		`- if (!result) {`
	`104`	`+ for (const pool of initialisedPools) {`
	`105`	`+ if (!pool) {`
	`106`	`+ // this means it was cancelled`
`106`	`107`	`return`
`107`	`108`	`}`
`108`	`109`
`109`		`- if (result.provider.mocker && result.provider.supportsParallelism) {`
`110`		`- parallelPools.push(result.runTests)`
	`110`	`+ if (pool.provider.mocker && pool.provider.supportsParallelism) {`
	`111`	`+ parallelPools.push(pool.runTests)`
`111`	`112`	`}`
`112`	`113`	`else {`
`113`		`- nonParallelPools.push(result.runTests)`
	`114`	`+ nonParallelPools.push(pool.runTests)`
`114`	`115`	`}`
`115`	`116`	`}`
`116`	`117`
Original file line number	Diff line number	Diff line change
`@@ -58,15 +58,17 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke`
`58`	`58`	`)`
`59`	`59`	`}`
`60`	`60`
`61`		`- if (!vitest._browserSessions.sessionIds.has(sessionId)) {`
`62`		`- const ids = [...vitest._browserSessions.sessionIds].join(', ')`
	`61`	`+ const sessions = vitest._browserSessions`
	`62`	`+`
	`63`	`+ if (!sessions.sessionIds.has(sessionId)) {`
	`64`	`+ const ids = [...sessions.sessionIds].join(', ')`
`63`	`65`	`return error(`
`64`	`66`	new Error(`[vitest] Unknown session id "${sessionId}". Expected one of ${ids}.`),
`65`	`67`	`)`
`66`	`68`	`}`
`67`	`69`
`68`	`70`	`if (type === 'orchestrator') {`
`69`		`- const session = vitest._browserSessions.getSession(sessionId)`
	`71`	`+ const session = sessions.getSession(sessionId)`
`70`	`72`	`// it's possible the session was already resolved by the preview provider`
`71`	`73`	`session?.connected()`
`72`	`74`	`}`
`@@ -94,7 +96,7 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke`
`94`	`96`	`clients.delete(rpcId)`
`95`	`97`	`globalServer.removeCDPHandler(rpcId)`
`96`	`98`	`if (type === 'orchestrator') {`
`97`		`- vitest._browserSessions.destroySession(sessionId)`
	`99`	`+ sessions.destroySession(sessionId)`
`98`	`100`	`}`
`99`	`101`	`// this will reject any hanging methods if there are any`
`100`	`102`	`rpc.$close(`