fix(webdriverio): properly switch context using switchToParentFrame() (webdriverio#14070)

christian-bromann · Delta456 · web-flow · commit 33a1b1967002 · 2025-01-17T10:45:03.000-08:00
* fix(webdriverio): properly switch context using switchToParentFrame() * fix e2e test * fix unit test * revert wrong change * fix(webdriverio): allow to take screenshots of frames (webdriverio#14071) * fix(webdriverio): allow to take screenshots of frames * Update packages/webdriverio/src/commands/browser/saveScreenshot.ts Co-authored-by: Swastik Baranwal <swstkbaranwal@gmail.com> --------- Co-authored-by: Swastik Baranwal <swstkbaranwal@gmail.com> * fix(webdriverio): properly switch context using switchToParentFrame() * pr feedback --------- Co-authored-by: Swastik Baranwal <swstkbaranwal@gmail.com>
diff --git a/.gitignore b/.gitignore
@@ -117,3 +117,4 @@ examples/wdio/mocha/*.json
 e2e/screenshot.png
 website/community/events
 website/events.json
+e2e/wdio/headless/*.png
diff --git a/__mocks__/fetch.ts b/__mocks__/fetch.ts
@@ -290,6 +290,8 @@ const requestMock: any = vi.fn().mockImplementation((uri, params) => {
             result = true
         } else if (body.script.includes('mobile:')) {
             result = true
+        } else if (body.script.includes('document.URL')) {
+            result = 'https://webdriver.io/?foo=bar'
         } else {
             result = script.apply(this, args)
         }
diff --git a/e2e/wdio/headless/test.e2e.ts b/e2e/wdio/headless/test.e2e.ts
@@ -5,6 +5,8 @@ import url from 'node:url'
 import path from 'node:path'
 import { browser, $, expect } from '@wdio/globals'
 
+import { imageSize } from 'image-size'
+
 const __dirname = path.dirname(url.fileURLToPath(import.meta.url))
 
 describe('main suite 1', () => {
@@ -549,6 +551,48 @@ describe('main suite 1', () => {
             await browser.switchFrame($('iframe'))
             await expect($('#tinymce')).toBePresent()
         })
+
+        describe('switchToParentFrame', () => {
+            it('switches to parent (not top-level)', async () => {
+                await browser.url('https://guinea-pig.webdriver.io/iframe.html')
+                await expect($('h1')).toHaveText('Frame Demo')
+                await expect($('h2')).not.toExist()
+                await expect($('h3')).not.toExist()
+
+                await browser.switchFrame($('#A'))
+                await expect($('h1')).not.toExist()
+                await expect($('h2')).toHaveText('IFrame A')
+                await expect($('h3')).not.toExist()
+
+                await browser.switchFrame($('#A2'))
+                await expect($('h1')).not.toExist()
+                await expect($('h2')).not.toExist()
+                await expect($('h3')).toHaveText('IFrame A2')
+
+                await browser.switchToParentFrame()
+                await expect($('h1')).not.toExist()
+                await expect($('h2')).toHaveText('IFrame A')
+                await expect($('h3')).not.toExist()
+            })
+
+            after(() => browser.switchFrame(null))
+        })
+
+        describe('taking screenshots', () => {
+            it('should take a screenshot of the iframe', async () => {
+                await browser.url('https://guinea-pig.webdriver.io/iframe.html')
+                await browser.switchFrame($('#A'))
+                await browser.switchFrame($('#A2'))
+
+                const screenshotPath = path.resolve(__dirname, 'iframe.png')
+                await browser.saveScreenshot(screenshotPath)
+                const dimensions = imageSize(screenshotPath)
+                expect(dimensions.width).toBe(187)
+                expect(dimensions.height).toBe(85)
+            })
+
+            after(() => browser.switchFrame(null))
+        })
     })
 
     describe('open resources with different protocols', () => {
diff --git a/package.json b/package.json
@@ -107,6 +107,7 @@
     "glob": "^11.0.0",
     "globals": "^15.12.0",
     "husky": "^9.1.4",
+    "image-size": "^1.2.0",
     "inquirer": "^11.0.1",
     "jsdom": "^25.0.0",
     "lerna": "8.1.9",
diff --git a/packages/webdriverio/src/commands/browser/switchFrame.ts b/packages/webdriverio/src/commands/browser/switchFrame.ts
@@ -1,14 +1,13 @@
 import logger from '@wdio/logger'
-import { ELEMENT_KEY, type local, type remote } from 'webdriver'
+import { ELEMENT_KEY, type remote } from 'webdriver'
 import type { ElementReference } from '@wdio/protocols'
 
-import { getContextManager } from '../../session/context.js'
+import { getContextManager, type FlatContextTree } from '../../session/context.js'
 import { LocalValue } from '../../utils/bidi/value.js'
 import { parseScriptResult } from '../../utils/bidi/index.js'
 import { SCRIPT_PREFIX, SCRIPT_SUFFIX } from '../constant.js'
 import type { ChainablePromiseElement } from '../../types.js'
 
-type FlatContextTree = Omit<local.BrowsingContextInfo, 'children'> & { children: string[] }
 const log = logger('webdriverio:switchFrame')
 
 /**
@@ -94,6 +93,8 @@ export async function switchFrame (
         return switchToFrame(this, context)
     }
 
+    const sessionContext = getContextManager(this)
+
     /**
      * if context is `null` the user is switching to the top level frame
      * which is always represented by the value of `getWindowHandle`
@@ -115,15 +116,15 @@ export async function switchFrame (
         let newContextId: string | undefined
 
         const urlContext = (
-            findContext(context, tree.contexts, byUrl) ||
+            sessionContext.findContext(context, tree.contexts, 'byUrl') ||
             /**
              * In case the user provides an url without `/` at the end, e.g. `https://example.com`,
              * the `browsingContextGetTree` command may return a context with the url `https://example.com/`.
              */
-            findContext(`${context}/`, tree.contexts, byUrl)
+            sessionContext.findContext(`${context}/`, tree.contexts, 'byUrl')
         )
-        const urlContextContaining = findContext(context, tree.contexts, byUrlContaining)
-        const contextIdContext = findContext(context, tree.contexts, byContextId)
+        const urlContextContaining = sessionContext.findContext(context, tree.contexts, 'byUrlContaining')
+        const contextIdContext = sessionContext.findContext(context, tree.contexts, 'byContextId')
         if (urlContext) {
             log.info(`Found context by url "${urlContext.url}" with context id "${urlContext.context}"`)
             newContextId = urlContext.context
@@ -139,9 +140,8 @@ export async function switchFrame (
             throw new Error(`No frame with url or id "${context}" found!`)
         }
 
-        const sessionContext = getContextManager(this)
         const currentContext = await sessionContext.getCurrentContext()
-        const allContexts = await getFlatContextTree(this)
+        const allContexts = await sessionContext.getFlatContextTree()
 
         /**
          * Fetch all iframes located in any available frame
@@ -252,6 +252,7 @@ export async function switchFrame (
             await switchToFrame(this, contextToSwitch.frameElement)
         }
 
+        sessionContext.setCurrentContext(newContextId)
         return newContextId
     }
 
@@ -273,7 +274,7 @@ export async function switchFrame (
      * the function for each of them.
      */
     if (typeof context === 'function') {
-        const allContexts = await getFlatContextTree(this)
+        const allContexts = await sessionContext.getFlatContextTree()
         const allContextIds = Object.keys(allContexts)
         for (const contextId of allContextIds) {
             const functionDeclaration = new Function(`
@@ -319,7 +320,6 @@ function switchToFrameHelper (browser: WebdriverIO.Browser, context: string) {
 }
 
 async function switchToFrameUsingElement (browser: WebdriverIO.Browser, element: WebdriverIO.Element) {
-    // await switchToFrame(browser, element)
     const frame = await browser.execute(
         (iframe: unknown) => (iframe as HTMLIFrameElement).contentWindow,
         element
@@ -332,60 +332,6 @@ async function switchToFrameUsingElement (browser: WebdriverIO.Browser, element:
     return frame.context
 }
 
-function byUrl (context: local.BrowsingContextInfo, url: string) {
-    return context.url === url
-}
-
-function byUrlContaining (context: local.BrowsingContextInfo, url: string) {
-    return context.url.includes(url)
-}
-
-function byContextId (context: local.BrowsingContextInfo, contextId: string) {
-    return context.context === contextId
-}
-
-function findContext (
-    urlOrId: string,
-    contexts: local.BrowsingContextInfoList | null,
-    matcher: typeof byUrl | typeof byUrlContaining | typeof byContextId
-): local.BrowsingContextInfo | undefined {
-    for (const context of contexts || []) {
-        if (matcher(context, urlOrId)) {
-            return context
-        }
-
-        if (Array.isArray(context.children) && context.children.length > 0) {
-            const result = findContext(urlOrId, context.children, matcher)
-            if (result) {
-                return result
-            }
-        }
-    }
-
-    return undefined
-}
-
-async function getFlatContextTree (browser: WebdriverIO.Browser): Promise<Record<string, FlatContextTree>> {
-    const tree = await browser.browsingContextGetTree({})
-
-    const mapContext = (context: local.BrowsingContextInfo): string[] => [
-        context.context,
-        ...(context.children || []).map(mapContext).flat(Infinity) as string[]
-    ]
-
-    /**
-     * transform context tree into a flat list of context objects with references
-     * to children
-     */
-    const allContexts: Record<string, FlatContextTree> = tree.contexts.map(mapContext).flat(Infinity)
-        .reduce((acc, ctx: string) => {
-            const context = findContext(ctx, tree.contexts, byContextId)
-            acc[ctx] = context as unknown as FlatContextTree
-            return acc
-        }, {} as Record<string, FlatContextTree>)
-    return allContexts
-}
-
 /**
  * While we deprecated the `switchToFrame` command for users, we still
  * have to use it internally to enable support for WebDriver Classic.
@@ -394,7 +340,16 @@ async function getFlatContextTree (browser: WebdriverIO.Browser): Promise<Record
  */
 function switchToFrame (browser: WebdriverIO.Browser, frame: ElementReference | number | null) {
     process.env.DISABLE_WEBDRIVERIO_DEPRECATION_WARNINGS = 'true'
-    return browser.switchToFrame(frame).finally(() => {
+    return browser.switchToFrame(frame).finally(async () => {
+        const sessionContext = getContextManager(browser)
+        const [frameTree, documentUrl] = await Promise.all([
+            sessionContext.getFlatContextTree(),
+            browser.execute(() => document.URL)
+        ])
+        const frame = Object.values(frameTree).find((ctx) => ctx.url === documentUrl)
+        if (frame) {
+            switchToFrameHelper(browser, frame.context)
+        }
         delete process.env.DISABLE_WEBDRIVERIO_DEPRECATION_WARNINGS
     })
 }
diff --git a/packages/webdriverio/src/node/saveScreenshot.ts b/packages/webdriverio/src/node/saveScreenshot.ts
@@ -49,7 +49,18 @@ export async function saveScreenshot (
         const browser = getBrowserObject(this)
         const contextManager = getContextManager(browser)
         const context = await contextManager.getCurrentContext()
-        const { data } = await this.browsingContextCaptureScreenshot({ context })
+        const tree = await this.browsingContextGetTree({})
+
+        /**
+         * WebDriver Bidi doesn't allow to take a screenshot of an iframe, it fails with:
+         * "unsupported operation - Non-top-level 'context'". Therefor we need to check if
+         * we are within an iframe and if so, take a screenshot of the document element
+         * instead.
+         */
+        const { data } = contextManager.findParentContext(context, tree.contexts)
+            ? await browser.$('html').getElement().then(
+                (el) => this.takeElementScreenshot(el.elementId).then((data) => ({ data })))
+            : await this.browsingContextCaptureScreenshot({ context })
         screenBuffer = data
     } else {
         screenBuffer = await this.takeScreenshot()
diff --git a/packages/webdriverio/src/session/context.ts b/packages/webdriverio/src/session/context.ts
@@ -1,3 +1,4 @@
+import type { local } from 'webdriver'
 import logger from '@wdio/logger'
 
 import { SessionManager } from './session.js'
@@ -10,6 +11,8 @@ export function getContextManager(browser: WebdriverIO.Browser) {
     return SessionManager.getSessionManager(browser, ContextManager)
 }
 
+export type FlatContextTree = Omit<local.BrowsingContextInfo, 'children'> & { children: string[] }
+
 /**
  * This class is responsible for managing context in a WebDriver session. Many BiDi commands
  * require to be executed in a specific context. This class is responsible for keeping track
@@ -81,6 +84,23 @@ export class ContextManager extends SessionManager {
     }
 
     #onCommand(event: { command: string, body: unknown }) {
+        /**
+         * update frame context if user switches using 'switchToParentFrame'
+         */
+        if (event.command === 'switchToParentFrame') {
+            if (!this.#currentContext) {
+                return
+            }
+
+            return this.#browser.browsingContextGetTree({}).then(({ contexts }) => {
+                const parentContext = this.findParentContext(this.#currentContext!, contexts)
+                if (!parentContext) {
+                    return
+                }
+                this.setCurrentContext(parentContext.context)
+            })
+        }
+
         /**
          * update frame context if user switches using 'switchToWindow'
          * which is WebDriver Classic only
@@ -181,4 +201,97 @@ export class ContextManager extends SessionManager {
     get mobileContext() {
         return this.#mobileContext
     }
+
+    /**
+     * Get the flat context tree for the current session
+     * @returns a flat list of all contexts in the current session
+     */
+    async getFlatContextTree (): Promise<Record<string, FlatContextTree>> {
+        const tree = await this.#browser.browsingContextGetTree({})
+
+        const mapContext = (context: local.BrowsingContextInfo): string[] => [
+            context.context,
+            ...(context.children || []).map(mapContext).flat(Infinity) as string[]
+        ]
+
+        /**
+         * transform context tree into a flat list of context objects with references
+         * to children
+         */
+        const allContexts: Record<string, FlatContextTree> = tree.contexts.map(mapContext).flat(Infinity)
+            .reduce((acc, ctx: string) => {
+                const context = this.findContext(ctx, tree.contexts, 'byContextId')
+                acc[ctx] = context as unknown as FlatContextTree
+                return acc
+            }, {} as Record<string, FlatContextTree>)
+        return allContexts
+    }
+
+    /**
+     * Find the parent context of a given context id
+     * @param contextId the context id you want to find the parent of
+     * @param contexts  the list of contexts to search through returned from `browsingContextGetTree`
+     * @returns         the parent context of the context with the given id
+     */
+    findParentContext (contextId: string, contexts: local.BrowsingContextInfoList): local.BrowsingContextInfo | undefined {
+        for (const context of contexts) {
+            if (context.children?.some((child) => child.context === contextId)) {
+                return context
+            }
+
+            if (Array.isArray(context.children) && context.children.length > 0) {
+                const result = this.findParentContext(contextId, context.children)
+                if (result) {
+                    return result
+                }
+            }
+        }
+
+        return undefined
+    }
+
+    /**
+     * Find a context by URL or ID
+     * @param urlOrId     The URL or ID of the context to find
+     * @param contexts    The list of contexts to search through returned from `browsingContextGetTree`
+     * @param matcherType The type of matcher to use to find the context
+     * @returns           The context with the given URL or ID
+     */
+    findContext (
+        urlOrId: string,
+        contexts: local.BrowsingContextInfoList | null,
+        matcherType: 'byUrl' | 'byUrlContaining' | 'byContextId'
+    ): local.BrowsingContextInfo | undefined {
+        const matcher = {
+            byUrl,
+            byUrlContaining,
+            byContextId
+        }[matcherType]
+        for (const context of contexts || []) {
+            if (matcher(context, urlOrId)) {
+                return context
+            }
+
+            if (Array.isArray(context.children) && context.children.length > 0) {
+                const result = this.findContext(urlOrId, context.children, matcherType)
+                if (result) {
+                    return result
+                }
+            }
+        }
+
+        return undefined
+    }
+}
+
+function byUrl (context: local.BrowsingContextInfo, url: string) {
+    return context.url === url
+}
+
+function byUrlContaining (context: local.BrowsingContextInfo, url: string) {
+    return context.url.includes(url)
+}
+
+function byContextId (context: local.BrowsingContextInfo, contextId: string) {
+    return context.context === contextId
 }
diff --git a/packages/webdriverio/tests/commands/browser/switchFrame.test.ts b/packages/webdriverio/tests/commands/browser/switchFrame.test.ts
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml

Original file line number	Diff line number	Diff line change
`@@ -290,6 +290,8 @@ const requestMock: any = vi.fn().mockImplementation((uri, params) => {`
`290`	`290`	`result = true`
`291`	`291`	`} else if (body.script.includes('mobile:')) {`
`292`	`292`	`result = true`
	`293`	`+ } else if (body.script.includes('document.URL')) {`
	`294`	`+ result = 'https://webdriver.io/?foo=bar'`
`293`	`295`	`} else {`
`294`	`296`	`result = script.apply(this, args)`
`295`	`297`	`}`