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

Author: sheremet-va

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:

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 || {},

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)
       }
     }
 

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
   }
 }

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(

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

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