feat(webdriverio): allow more options for screenshot taking with Bidi (#14267)

* feat(webdriverio): allow more options for screenshot taking with Bidi

* simplify options

* fix unit tests

packages/webdriverio/src/commands/browser/saveScreenshot.ts

@@ -1,4 +1,5 @@
 import { environment } from '../../environment.js'
+import type { SaveScreenshotOptions } from '../../types.js'
 
 /**
  *
@@ -8,9 +9,25 @@ import { environment } from '../../environment.js'
  *
  * <example>
     :saveScreenshot.js
-    it('should save a screenshot of the browser view', async () => {
+    it('should save a screenshot of the browser viewport', async () => {
         await browser.saveScreenshot('./some/path/screenshot.png');
     });
+
+    it('should save a screenshot of the full page', async () => {
+        await browser.saveScreenshot('./some/path/screenshot.png', { fullPage: true });
+    });
+
+    it('should save a screenshot of a specific rectangle', async () => {
+        await browser.saveScreenshot('./some/path/screenshot.png', { clip: { x: 0, y: 0, width: 100, height: 100 } });
+    });
+
+    it('should save a screenshot of the full page in JPEG format', async () => {
+        await browser.saveScreenshot('./some/path/screenshot.jpeg', { fullPage: true, format: 'jpeg' });
+    });
+
+    it('should save a screenshot of the full page in JPEG format with quality 50', async () => {
+        await browser.saveScreenshot('./some/path/screenshot.jpeg', { fullPage: true, format: 'jpeg', quality: 50 });
+    });
  * </example>
  *
  * When running from a hook, make sure to explicitly define the hook as async:
@@ -22,16 +39,22 @@ import { environment } from '../../environment.js'
  * </example>
  * @alias browser.saveScreenshot
  * @param   {String}  filepath  path to the generated image (`.png` suffix is required) relative to the execution directory
- * @return  {Buffer}            screenshot buffer
+ * @param   {Object}  options   screenshot options
+ * @param   {Boolean} [options.fullPage=false]  whether to take a screenshot of the full page or just the current viewport
+ * @param   {String}  [options.format='png']    the format of the screenshot (either `png` or `jpeg`)
+ * @param   {Number}  [options.quality=100]     the quality of the screenshot in case of JPEG format in range 0-100 percent
+ * @param   {Object}  [options.clip]            clipping a rectangle of the screenshot
+ * @return  {Buffer}                            screenshot buffer
  * @type utility
  *
  */
 export async function saveScreenshot (
     this: WebdriverIO.Browser,
-    filepath: string
+    filepath: string,
+    options?: SaveScreenshotOptions
 ) {
     /**
      * run command implementation based on given environment
      */
-    return environment.value.saveScreenshot.call(this, filepath)
+    return environment.value.saveScreenshot.call(this, filepath, options)
 }

packages/webdriverio/src/node/saveScreenshot.ts

@@ -1,9 +1,11 @@
-import fs from 'node:fs'
-import { getBrowserObject } from '@wdio/utils'
+import fs from 'node:fs/promises'
+import path from 'node:path'
 
+import { getBrowserObject } from '@wdio/utils'
+import type { remote } from 'webdriver'
 import { getAbsoluteFilepath, assertDirectoryExists } from './utils.js'
 import { getContextManager } from '../session/context.js'
-
+import type { SaveScreenshotOptions } from '../types.js'
 /**
  *
  * Save a screenshot of the current browsing context to a PNG file on your OS. Be aware that
@@ -32,41 +34,110 @@ import { getContextManager } from '../session/context.js'
  */
 export async function saveScreenshot (
     this: WebdriverIO.Browser,
-    filepath: string
+    filepath: string,
+    options?: SaveScreenshotOptions
 ) {
     /**
      * type check
      */
-    if (typeof filepath !== 'string' || !filepath.endsWith('.png')) {
+    if (typeof filepath !== 'string') {
         throw new Error('saveScreenshot expects a filepath of type string and ".png" file ending')
     }
 
     const absoluteFilepath = getAbsoluteFilepath(filepath)
     await assertDirectoryExists(absoluteFilepath)
 
-    let screenBuffer: string
-    if (this.isBidi) {
-        const browser = getBrowserObject(this)
-        const contextManager = getContextManager(browser)
-        const context = await contextManager.getCurrentContext()
-        const tree = await this.browsingContextGetTree({})
+    const screenBuffer = this.isBidi
+        ? await takeScreenshotBidi.call(this, filepath, options)
+        : await takeScreenshotClassic.call(this, filepath, options)
 
-        /**
-         * 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()
-    }
     const screenshot = Buffer.from(screenBuffer, 'base64')
-    fs.writeFileSync(absoluteFilepath, screenshot)
+    await fs.writeFile(absoluteFilepath, screenshot)
 
     return screenshot
 }
+
+/**
+ * take screenshot using legacy WebDriver command
+ * @returns {string} a base64 encoded screenshot
+ */
+export function takeScreenshotClassic (this: WebdriverIO.Browser, filepath: string, options?: SaveScreenshotOptions): Promise<string> {
+    if (options) {
+        throw new Error('saveScreenshot does not support options in WebDriver Classic mode')
+    }
+    const fileExtension = path.extname(filepath).slice(1)
+    if (fileExtension !== 'png') {
+        throw new Error('Invalid file extension, use ".png" for PNG format')
+    }
+    return this.takeScreenshot()
+}
+
+/**
+ * takeScreenshotBidi
+ * @returns {string} a base64 encoded screenshot
+ */
+export async function takeScreenshotBidi (this: WebdriverIO.Browser, filepath: string, options?: SaveScreenshotOptions): Promise<string> {
+    const browser = getBrowserObject(this)
+    const contextManager = getContextManager(browser)
+    const context = await contextManager.getCurrentContext()
+    const tree = await this.browsingContextGetTree({})
+    const origin: remote.BrowsingContextCaptureScreenshotParameters['origin'] = options?.fullPage ? 'document' : 'viewport'
+    const givenFormat = options?.format || path.extname(filepath).slice(1)
+    const imageFormat = givenFormat === 'png'
+        ? 'image/png'
+        : givenFormat === 'jpeg' || givenFormat === 'jpg'
+            ? 'image/jpeg'
+            : undefined
+
+    if (!imageFormat) {
+        throw new Error(`Invalid image format, use 'png', 'jpg' or 'jpeg', got '${options?.format}'`)
+    }
+
+    if (imageFormat === 'image/jpeg' && path.extname(filepath) !== '.jpeg' && path.extname(filepath) !== '.jpg') {
+        throw new Error('Invalid file extension, use ".jpeg" or ".jpg" for JPEG format')
+    } else if (imageFormat === 'image/png' && path.extname(filepath) !== '.png') {
+        throw new Error('Invalid file extension, use ".png" for PNG format')
+    }
+
+    const quality = typeof options?.quality === 'number' ? (options.quality / 100) : undefined
+    if (typeof options?.quality === 'number' && (options?.quality < 0 || options?.quality > 100)) {
+        throw new Error(`Invalid quality, use a number between 0 and 100, got '${options?.quality}'`)
+    }
+
+    if (typeof options?.quality === 'number' && imageFormat !== 'image/jpeg') {
+        throw new Error('Invalid option "quality" for PNG format')
+    }
+
+    const format: remote.BrowsingContextImageFormat = {
+        type: imageFormat,
+        quality
+    }
+
+    const clip: remote.BrowsingContextBoxClipRectangle | undefined = options?.clip
+        ? {
+            type: 'box',
+            x: options.clip.x,
+            y: options.clip.y,
+            width: options.clip.width,
+            height: options.clip.height
+        }
+        : undefined
+    if (clip) {
+        if (typeof clip.x !== 'number' || typeof clip.y !== 'number' || typeof clip.width !== 'number' || typeof clip.height !== 'number') {
+            throw new Error('Invalid clip, use an object with x, y, width and height properties')
+        }
+    }
+
+    /**
+     * 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, origin, format, clip })
+    return data
+}
+

packages/webdriverio/src/types.ts

@@ -650,6 +650,33 @@ export interface ExtendedElementReference {
 export type SupportedScopes = 'geolocation' | 'userAgent' | 'colorScheme' | 'onLine' | 'clock' | 'device'
 export type RestoreMap = Map<SupportedScopes, (() => Promise<any>)[]>
 
+export interface SaveScreenshotOptions {
+    /**
+     * Whether to take a screenshot of the full page or just the current viewport.
+     * @default false
+     */
+    fullPage?: boolean
+    /**
+     * The format of the screenshot.
+     * @default 'png'
+     */
+    format?: 'png' | 'jpeg' | 'jpg'
+    /**
+     * The quality of the screenshot in case of JPEG format in range 0-100 percent.
+     * @default 100
+     */
+    quality?: number
+    /**
+     * Clipping a rectangle of the screenshot.
+     */
+    clip?: {
+        x: number
+        y: number
+        width: number
+        height: number
+    }
+}
+
 declare global {
     namespace WebdriverIO {
         /**

packages/webdriverio/tests/commands/browser/saveScreenshot.test.ts

@@ -1,5 +1,5 @@
 import { expect, describe, beforeEach, afterEach, it, vi, type MockInstance } from 'vitest'
-import fs from 'node:fs'
+import fs from 'node:fs/promises'
 import path from 'node:path'
 
 import { remote } from '../../../src/index.js'
@@ -7,7 +7,13 @@ import * as utils from '../../../src/node/utils.js'
 
 import '../../../src/node.js'
 
-vi.mock('fs')
+vi.mock('fs/promises', () => ({
+    default: {
+        writeFile: vi.fn().mockResolvedValue({}),
+        access: vi.fn().mockResolvedValue({})
+    }
+
+}))
 vi.mock('fetch')
 vi.mock('@wdio/logger', () => import(path.join(process.cwd(), '__mocks__', '@wdio/logger')))
 
@@ -15,7 +21,7 @@ describe('saveScreenshot', () => {
     let browser: WebdriverIO.Browser
     let getAbsoluteFilepathSpy: MockInstance
     let assertDirectoryExistsSpy: MockInstance
-    let writeFileSyncSpy: MockInstance
+    let writeFileSpy: MockInstance
 
     beforeEach(async () => {
         browser = await remote({
@@ -26,13 +32,13 @@ describe('saveScreenshot', () => {
         })
         getAbsoluteFilepathSpy = vi.spyOn(utils, 'getAbsoluteFilepath')
         assertDirectoryExistsSpy = vi.spyOn(utils, 'assertDirectoryExists')
-        writeFileSyncSpy = vi.spyOn(fs, 'writeFileSync')
+        writeFileSpy = vi.spyOn(fs, 'writeFile')
     })
 
     afterEach(() => {
         getAbsoluteFilepathSpy.mockClear()
         assertDirectoryExistsSpy.mockClear()
-        writeFileSyncSpy.mockClear()
+        writeFileSpy.mockClear()
     })
 
     it('should take screenshot of page', async () => {
@@ -54,22 +60,22 @@ describe('saveScreenshot', () => {
         expect(screenshot.toString()).toBe('some screenshot')
 
         // write to file
-        expect(writeFileSyncSpy).toHaveBeenCalledTimes(1)
-        expect(writeFileSyncSpy).toHaveBeenCalledWith(getAbsoluteFilepathSpy.mock.results[0].value, expect.any(Buffer))
+        expect(writeFileSpy).toHaveBeenCalledTimes(1)
+        expect(writeFileSpy).toHaveBeenCalledWith(getAbsoluteFilepathSpy.mock.results[0].value, expect.any(Buffer))
     })
 
     it('should fail if no filename provided', async () => {
-        const expectedError = new Error('saveScreenshot expects a filepath of type string and ".png" file ending')
+        const expectedError =
 
         // no file
         await expect(
             // @ts-ignore test invalid parameter
             browser.saveScreenshot()
-        ).rejects.toEqual(expectedError)
+        ).rejects.toEqual(new Error('saveScreenshot expects a filepath of type string and ".png" file ending'))
 
         // wrong extension
         await expect(
             browser.saveScreenshot('./file.txt')
-        ).rejects.toEqual(expectedError)
+        ).rejects.toEqual(new Error('Invalid file extension, use ".png" for PNG format'))
     })
 })

packages/webdriverio/tests/node/saveScreenshot.test.ts

@@ -0,0 +1,115 @@
+import { describe, it, expect, vi } from 'vitest'
+
+import { saveScreenshot } from '../../src/node/saveScreenshot.js'
+import { getContextManager } from '../../src/session/context.js'
+
+vi.mock('node:fs/promises', () => ({
+    default: {
+        access: vi.fn().mockResolvedValue(undefined),
+        writeFile: vi.fn()
+    }
+}))
+
+vi.mock('../../src/session/context.js', () => ({
+    getContextManager: vi.fn().mockReturnValue({
+        getCurrentContext: vi.fn().mockResolvedValue('context'),
+        findParentContext: vi.fn().mockResolvedValue({ data: 'bidi screenshot' })
+    })
+}))
+
+const browserClassic = {
+    isBidi: false,
+    takeScreenshot: vi.fn().mockResolvedValue(btoa('classic screenshot'))
+} as unknown as WebdriverIO.Browser
+
+const browserBidi = {
+    isBidi: true,
+    $: vi.fn().mockReturnValue({
+        getElement: vi.fn().mockResolvedValue({ elementId: 'elementId' })
+    }),
+    takeElementScreenshot: vi.fn().mockResolvedValue(btoa('bidi iframe screenshot')),
+    browsingContextGetTree: vi.fn().mockResolvedValue({ contexts: [] }),
+    browsingContextCaptureScreenshot: vi.fn().mockResolvedValue({ data: btoa('bidi screenshot') })
+} as unknown as WebdriverIO.Browser
+
+describe('saveScreenshot', () => {
+    describe('WebDriver Classic mode', () => {
+        it('should save a screenshot of the browser viewport', async () => {
+            const screenshot = await saveScreenshot.call(browserClassic, 'screenshot.png')
+            expect(screenshot.toString()).toBe('classic screenshot')
+        })
+
+        it('should fail if options are provided', async () => {
+            await expect(saveScreenshot.call(browserClassic, 'screenshot.png', { fullPage: true })).rejects.toThrow()
+        })
+    })
+
+    describe('WebDriver Bidi mode', () => {
+        it('should save a screenshot of an iframe', async () => {
+            const screenshot = await saveScreenshot.call(browserBidi, 'screenshot.png')
+            expect(screenshot.toString()).toBe('bidi iframe screenshot')
+        })
+
+        it('should save a screenshot of the full page', async () => {
+            const contextManager = getContextManager(browserBidi)
+            vi.mocked(contextManager.findParentContext).mockReturnValue(undefined)
+            const screenshot = await saveScreenshot.call(browserBidi, 'screenshot.png', { fullPage: true })
+            expect(screenshot.toString()).toBe('bidi screenshot')
+            expect(browserBidi.browsingContextCaptureScreenshot).toHaveBeenCalledWith({
+                context: 'context',
+                origin: 'document',
+                format: {
+                    type: 'image/png',
+                    quality: undefined
+                },
+                clip: undefined
+            })
+        })
+
+        it('should save a screenshot of a specific rectangle', async () => {
+            const screenshot = await saveScreenshot.call(browserBidi, 'screenshot.png', { clip: { x: 10, y: 10, width: 100, height: 100 } })
+            expect(screenshot.toString()).toBe('bidi screenshot')
+            expect(browserBidi.browsingContextCaptureScreenshot).toHaveBeenCalledWith({
+                context: 'context',
+                origin: 'viewport',
+                clip: {
+                    type: 'box',
+                    x: 10,
+                    y: 10,
+                    width: 100,
+                    height: 100
+                },
+                format: {
+                    type: 'image/png',
+                    quality: undefined
+                }
+            })
+        })
+
+        it('should throw an error if the image format is invalid', async () => {
+            // @ts-expect-error invalid format
+            await expect(saveScreenshot.call(browserBidi, 'screenshot.png', { format: 'invalid' })).rejects.toThrow()
+        })
+
+        it('should throw an error if the file extension is invalid', async () => {
+            await expect(saveScreenshot.call(browserBidi, 'screenshot.jpeg', { format: 'png' })).rejects.toThrow()
+        })
+
+        it('should throw an error if the file extension is invalid for JPEG format', async () => {
+            await saveScreenshot.call(browserBidi, 'screenshot.jpeg')
+            await expect(saveScreenshot.call(browserBidi, 'screenshot.exe', { format: 'jpeg' })).rejects.toThrow()
+        })
+
+        it('should throw an error if the quality is invalid', async () => {
+            await expect(saveScreenshot.call(browserBidi, 'screenshot.png', { quality: 101 }))
+                .rejects.toThrow(/Invalid quality, use a number between 0 and 100,/)
+            await expect(saveScreenshot.call(browserBidi, 'screenshot.png', { quality: 50 }))
+                .rejects.toThrow(/Invalid option "quality" for PNG format/)
+        })
+
+        it('should throw an error if the clip is invalid', async () => {
+            // @ts-expect-error invalid clip
+            await expect(saveScreenshot.call(browserBidi, 'screenshot.png', { clip: { x: 'invalid' } })).rejects.toThrow()
+        })
+    })
+})

scripts/changelog.ts

@@ -14,7 +14,7 @@ import { highlight } from 'cli-highlight'
 import { Changelog } from 'lerna-changelog'
 import { load } from 'lerna-changelog/lib/configuration.js'
 
-import pkg from '../lerna.json' assert { type: 'json' }
+import pkg from '../lerna.json' with { type: 'json' }
 
 const __dirname = path.dirname(url.fileURLToPath(import.meta.url))
 const root = path.resolve(__dirname, '..')