perf(browser): do wdio context switching only once per file

packages/browser/src/client/tester/context.ts

@@ -1,5 +1,4 @@
 import type { Options as TestingLibraryOptions, UserEvent as TestingLibraryUserEvent } from '@testing-library/user-event'
-import type { BrowserRPC } from '@vitest/browser/client'
 import type { RunnerTask } from 'vitest'
 import type {
   BrowserPage,
@@ -19,15 +18,11 @@ import { convertElementToCssSelector, ensureAwaited, getBrowserState, getWorkerS
 const state = () => getWorkerState()
 // @ts-expect-error not typed global
 const provider = __vitest_browser_runner__.provider
-function filepath() {
-  return getWorkerState().filepath || getWorkerState().current?.file?.filepath || undefined
-}
-const rpc = () => getWorkerState().rpc as any as BrowserRPC
 const sessionId = getBrowserState().sessionId
 const channel = new BroadcastChannel(`vitest:${sessionId}`)
 
 function triggerCommand<T>(command: string, ...args: any[]) {
-  return rpc().triggerCommand<T>(sessionId, command, filepath(), args)
+  return getBrowserState().commands.triggerCommand<T>(command, args)
 }
 
 export function createUserEvent(__tl_user_event_base__?: TestingLibraryUserEvent, options?: TestingLibraryOptions): UserEvent {
@@ -52,6 +47,10 @@ export function createUserEvent(__tl_user_event_base__?: TestingLibraryUserEvent
       return createUserEvent()
     },
     async cleanup() {
+      // avoid cleanup rpc call if there is nothing to cleanup
+      if (!keyboard.unreleased.length) {
+        return
+      }
       return ensureAwaited(async () => {
         await triggerCommand('__vitest_cleanup', keyboard)
         keyboard.unreleased = []
@@ -106,9 +105,7 @@ export function createUserEvent(__tl_user_event_base__?: TestingLibraryUserEvent
       })
     },
     tab(options: UserEventTabOptions = {}) {
-      return ensureAwaited(() => {
-        return triggerCommand('__vitest_tab', options)
-      })
+      return ensureAwaited(() => triggerCommand('__vitest_tab', options))
     },
     async keyboard(text: string) {
       return ensureAwaited(async () => {

packages/browser/src/client/tester/locators/index.ts

@@ -1,4 +1,3 @@
-import type { BrowserRPC } from '@vitest/browser/client'
 import type {
   LocatorByRoleOptions,
   LocatorOptions,
@@ -8,8 +7,6 @@ import type {
   UserEventFillOptions,
   UserEventHoverOptions,
 } from '@vitest/browser/context'
-import type { WorkerGlobalState } from 'vitest'
-import type { BrowserRunnerState } from '../../utils'
 import { page, server } from '@vitest/browser/context'
 import {
   getByAltTextSelector,
@@ -22,7 +19,7 @@ import {
   Ivya,
   type ParsedSelector,
 } from 'ivya'
-import { ensureAwaited, getBrowserState, getWorkerState } from '../../utils'
+import { ensureAwaited, getBrowserState } from '../../utils'
 import { getElementError } from '../public-utils'
 
 // we prefer using playwright locators because they are more powerful and support Shadow DOM
@@ -205,27 +202,10 @@ export abstract class Locator {
     return this.selector
   }
 
-  private get state(): BrowserRunnerState {
-    return getBrowserState()
-  }
-
-  private get worker(): WorkerGlobalState {
-    return getWorkerState()
-  }
-
-  private get rpc(): BrowserRPC {
-    return this.worker.rpc as any as BrowserRPC
-  }
-
   protected triggerCommand<T>(command: string, ...args: any[]): Promise<T> {
-    const filepath = this.worker.filepath
-      || this.worker.current?.file?.filepath
-      || undefined
-
-    return ensureAwaited(() => this.rpc.triggerCommand<T>(
-      this.state.sessionId,
+    const commands = getBrowserState().commands
+    return ensureAwaited(() => commands.triggerCommand<T>(
       command,
-      filepath,
       args,
     ))
   }

packages/browser/src/client/tester/runner.ts

@@ -1,19 +1,21 @@
 import type { CancelReason, File, Suite, Task, TaskEventPack, TaskResultPack, VitestRunner } from '@vitest/runner'
 import type { SerializedConfig, WorkerGlobalState } from 'vitest'
 import type { VitestExecutor } from 'vitest/execute'
+import type { CommandsManager } from '../utils'
 import type { VitestBrowserClientMocker } from './mocker'
 import { globalChannel } from '@vitest/browser/client'
-import { page, userEvent } from '@vitest/browser/context'
+import { page, server, userEvent } from '@vitest/browser/context'
 import { loadDiffConfig, loadSnapshotSerializers, takeCoverageInsideWorker } from 'vitest/browser'
 import { NodeBenchmarkRunner, VitestTestRunner } from 'vitest/runners'
 import { originalPositionFor, TraceMap } from 'vitest/utils'
 import { createStackString, parseStacktrace } from '../../../../utils/src/source-map'
-import { executor, getWorkerState } from '../utils'
+import { executor, getBrowserState, getWorkerState } from '../utils'
 import { rpc } from './rpc'
 import { VitestBrowserSnapshotEnvironment } from './snapshot'
 
 interface BrowserRunnerOptions {
   config: SerializedConfig
+  commands: CommandsManager
 }
 
 export const browserHashMap: Map<string, string> = new Map()
@@ -32,10 +34,28 @@ export function createBrowserRunner(
     public config: SerializedConfig
     hashMap = browserHashMap
     public sourceMapCache = new Map<string, any>()
+    private commands: CommandsManager
+
+    private firstCommand = true
+    private contextSwitched = false
 
     constructor(options: BrowserRunnerOptions) {
       super(options.config)
       this.config = options.config
+      this.commands = options.commands
+
+      // webdiverio context depends on the iframe state, so we need to switch the context
+      // for every test file
+      if (server.provider === 'webdriverio') {
+        this.commands.onCommand(async () => {
+          // if this is the first command, make sure we switched the command context to an iframe
+          if (this.firstCommand) {
+            this.firstCommand = false
+            await rpc().wdioSwitchContext('iframe')
+            this.contextSwitched = true
+          }
+        })
+      }
     }
 
     onBeforeTryTask: VitestRunner['onBeforeTryTask'] = async (...args) => {
@@ -89,11 +109,18 @@ export function createBrowserRunner(
       ])
     }
 
+    // this always has a single test file in Vitest
     onAfterRunFiles = async (files: File[]) => {
       const [coverage] = await Promise.all([
         coverageModule?.takeCoverage?.(),
         mocker.invalidate(),
         super.onAfterRunFiles?.(files),
+        (async () => {
+          if (this.contextSwitched) {
+            this.contextSwitched = false
+            await rpc().wdioSwitchContext('parent')
+          }
+        })(),
       ])
 
       if (coverage) {
@@ -170,6 +197,7 @@ export async function initiateRunner(
   }
   const runner = new BrowserRunner({
     config,
+    commands: getBrowserState().commands,
   })
 
   const [diffOptions] = await Promise.all([

packages/browser/src/client/tester/tester.ts

@@ -1,7 +1,7 @@
 import { channel, client, onCancel } from '@vitest/browser/client'
 import { page, userEvent } from '@vitest/browser/context'
 import { collectTests, setupCommonEnv, SpyModule, startCoverageInsideWorker, startTests, stopCoverageInsideWorker } from 'vitest/browser'
-import { executor, getBrowserState, getConfig, getWorkerState } from '../utils'
+import { CommandsManager, executor, getBrowserState, getConfig, getWorkerState } from '../utils'
 import { setupDialogsSpy } from './dialog'
 import { setupExpectDom } from './expect-element'
 import { setupConsoleLogSpy } from './logger'
@@ -34,6 +34,8 @@ async function prepareTestEnvironment(files: string[]) {
   state.onCancel = onCancel
   state.rpc = rpc as any
 
+  getBrowserState().commands = new CommandsManager()
+
   // TODO: expose `worker`
   const interceptor = createModuleMockerInterceptor()
   const mocker = new VitestBrowserClientMocker(

packages/browser/src/client/utils.ts

@@ -1,4 +1,5 @@
 import type { SerializedConfig, WorkerGlobalState } from 'vitest'
+import type { BrowserRPC } from './client'
 
 export async function importId(id: string): Promise<any> {
   const name = `/@id/${id}`.replace(/\\/g, '/')
@@ -77,6 +78,7 @@ export interface BrowserRunnerState {
   method: 'run' | 'collect'
   runTests?: (tests: string[]) => Promise<void>
   createTesters?: (files: string[]) => Promise<void>
+  commands: CommandsManager
   cdp?: {
     on: (event: string, listener: (payload: any) => void) => void
     once: (event: string, listener: (payload: any) => void) => void
@@ -194,3 +196,28 @@ function getParent(el: Element) {
   }
   return parent
 }
+
+export class CommandsManager {
+  public history: string[] = []
+  private _listeners: ((command: string, args: any[]) => void)[] = []
+
+  public onCommand(listener: (command: string, args: any[]) => void): void {
+    this._listeners.push(listener)
+  }
+
+  public async triggerCommand<T>(command: string, args: any[]): Promise<T> {
+    this.history.push(command)
+    const state = getWorkerState()
+    const rpc = state.rpc as any as BrowserRPC
+    const { sessionId } = getBrowserState()
+    const filepath = state.filepath || state.current?.file?.filepath
+    if (this._listeners.length) {
+      await Promise.all(this._listeners.map(listener => listener(command, args)))
+    }
+    return rpc.triggerCommand<T>(sessionId, command, filepath, args)
+  }
+
+  public reset(): void {
+    this.history = []
+  }
+}

packages/browser/src/node/commands/keyboard.ts

@@ -54,6 +54,9 @@ export const keyboardCleanup: UserEventCommand<(state: KeyboardState) => Promise
   state,
 ) => {
   const { provider, sessionId } = context
+  if (!state.unreleased) {
+    return
+  }
   if (provider instanceof PlaywrightBrowserProvider) {
     const page = provider.getPage(sessionId)
     for (const key of state.unreleased) {

packages/browser/src/node/plugins/pluginContext.ts

@@ -32,15 +32,13 @@ async function generateContextFile(
   globalServer: ParentBrowserProject,
 ) {
   const commands = Object.keys(globalServer.commands)
-  const filepathCode
-    = '__vitest_worker__.filepath || __vitest_worker__.current?.file?.filepath || undefined'
   const provider = [...globalServer.children][0].provider || { name: 'preview' }
   const providerName = provider.name
 
   const commandsCode = commands
     .filter(command => !command.startsWith('__vitest'))
     .map((command) => {
-      return `    ["${command}"]: (...args) => rpc().triggerCommand(sessionId, "${command}", filepath(), args),`
+      return `    ["${command}"]: (...args) => __vitest_browser_runner__.commands.triggerCommand("${command}", args),`
     })
     .join('\n')
 
@@ -53,9 +51,6 @@ async function generateContextFile(
   return `
 import { page, createUserEvent, cdp } from '${distContextPath}'
 ${userEventNonProviderImport}
-const filepath = () => ${filepathCode}
-const rpc = () => __vitest_worker__.rpc
-const sessionId = __vitest_browser_runner__.sessionId
 
 export const server = {
   platform: ${JSON.stringify(process.platform)},

packages/browser/src/node/providers/webdriver.ts

@@ -37,7 +37,7 @@ export class WebdriverBrowserProvider implements BrowserProvider {
     this.options = options as RemoteOptions
   }
 
-  async beforeCommand(): Promise<void> {
+  async switchToTestFrame(): Promise<void> {
     const page = this.browser!
     const iframe = await page.findElement(
       'css selector',
@@ -46,7 +46,7 @@ export class WebdriverBrowserProvider implements BrowserProvider {
     await page.switchToFrame(iframe)
   }
 
-  async afterCommand(): Promise<void> {
+  async switchToMainFrame(): Promise<void> {
     await this.browser!.switchToParentFrame()
   }
 

packages/browser/src/node/rpc.ts

@@ -12,6 +12,7 @@ import { parse, stringify } from 'flatted'
 import { dirname } from 'pathe'
 import { createDebugger, isFileServingAllowed, isValidApiRequest } from 'vitest/node'
 import { WebSocketServer } from 'ws'
+import { WebdriverBrowserProvider } from './providers/webdriver'
 
 const debug = createDebugger('vitest:browser:api')
 
@@ -203,6 +204,20 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject): void {
         getCountOfFailedTests() {
           return vitest.state.getCountOfFailedTests()
         },
+        async wdioSwitchContext(direction) {
+          const provider = project.browser!.provider as WebdriverBrowserProvider
+          if (!provider) {
+            throw new Error('Commands are only available for browser tests.')
+          }
+          if (provider.name !== 'webdriverio') {
+            throw new Error('Switch context is only available for WebDriverIO provider.')
+          }
+          if (direction === 'iframe') {
+            await provider.switchToTestFrame()
+          } else {
+            await provider.switchToMainFrame()
+          }
+        },
         async triggerCommand(sessionId, command, testPath, payload) {
           debug?.('[%s] Triggering command "%s"', sessionId, command)
           const provider = project.browser!.provider
@@ -213,7 +228,6 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject): void {
           if (!commands || !commands[command]) {
             throw new Error(`Unknown command "${command}".`)
           }
-          await provider.beforeCommand?.(command, payload)
           const context = Object.assign(
             {
               testPath,
@@ -224,14 +238,7 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject): void {
             },
             provider.getCommandsContext(sessionId),
           ) as any as BrowserCommandContext
-          let result
-          try {
-            result = await commands[command](context, ...payload)
-          }
-          finally {
-            await provider.afterCommand?.(command, payload)
-          }
-          return result
+          return await commands[command](context, ...payload)
         },
         finishBrowserTests(sessionId: string) {
           debug?.('[%s] Finishing browser tests for session', sessionId)

packages/browser/src/node/types.ts

@@ -39,6 +39,7 @@ export interface WebSocketBrowserHandlers {
   getBrowserFileSourceMap: (
     id: string
   ) => SourceMap | null | { mappings: '' } | undefined
+  wdioSwitchContext: (direction: 'iframe' | 'parent') => void
 
   // cdp
   sendCdpEvent: (sessionId: string, event: string, payload?: Record<string, unknown>) => unknown

packages/vitest/src/node/types/browser.ts

@@ -24,8 +24,6 @@ export interface BrowserProvider {
    */
   supportsParallelism: boolean
   getSupportedBrowsers: () => readonly string[]
-  beforeCommand?: (command: string, args: unknown[]) => Awaitable<void>
-  afterCommand?: (command: string, args: unknown[]) => Awaitable<void>
   getCommandsContext: (sessionId: string) => Record<string, unknown>
   openPage: (sessionId: string, url: string, beforeNavigate?: () => Promise<void>) => Promise<void>
   getCDPSession?: (sessionId: string) => Promise<CDPSession>