feat(browser): add screenshot.save option (#7777)

Author: sheremet-va

docs/guide/browser/context.md

@@ -105,6 +105,11 @@ export const page: {
 The `getBy*` API is explained at [Locators API](/guide/browser/locators).
 :::
 
+::: warning WARNING <Version>3.1.2</Version>
+Note that `screenshot` will always return a base64 string if `save` is set to `false`.
+The `path` is also ignored in that case.
+:::
+
 ## `cdp`
 
 The `cdp` export returns the current Chrome DevTools Protocol session. It is mostly useful to library authors to build tools on top of it.

docs/guide/browser/locators.md

@@ -769,6 +769,7 @@ await languages.selectOptions([
 ### screenshot
 
 ```ts
+function screenshot(options: LocatorScreenshotOptions & { save: false }): Promise<string>
 function screenshot(options: LocatorScreenshotOptions & { base64: true }): Promise<{
   path: string
   base64: string
@@ -797,6 +798,11 @@ const { path, base64 } = await button.screenshot({
 // bas64 - base64 encoded string of the screenshot
 ```
 
+::: warning WARNING <Version>3.1.2</Version>
+Note that `screenshot` will always return a base64 string if `save` is set to `false`.
+The `path` is also ignored in that case.
+:::
+
 ### query
 
 ```ts

packages/browser/context.d.ts

@@ -29,12 +29,19 @@ export interface ScreenshotOptions {
   element?: Element | Locator
   /**
    * Path relative to the current test file.
+   * @default `__screenshots__/${testFileName}/${testName}.png`
    */
   path?: string
   /**
    * Will also return the base64 encoded screenshot alongside the path.
    */
   base64?: boolean
+  /**
+   * Keep the screenshot on the file system. If file is not saved,
+   * `page.screenshot` always returns `base64` screenshot.
+   * @default true
+   */
+  save?: boolean
 }
 
 export interface BrowserCommands {
@@ -552,11 +559,16 @@ export interface BrowserPage extends LocatorSelectors {
    * Make a screenshot of the test iframe or a specific element.
    * @returns Path to the screenshot file or path and base64.
    */
+  screenshot(options: Omit<ScreenshotOptions, 'save'> & { save: false }): Promise<string>
   screenshot(options: Omit<ScreenshotOptions, 'base64'> & { base64: true }): Promise<{
     path: string
     base64: string
   }>
-  screenshot(options?: ScreenshotOptions): Promise<string>
+  screenshot(options?: Omit<ScreenshotOptions, 'base64'>): Promise<string>
+  screenshot(options?: ScreenshotOptions): Promise<string | {
+    path: string
+    base64: string
+  }>
   /**
    * Extend default `page` object with custom methods.
    */

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

@@ -1,6 +1,6 @@
 import type { BrowserCommand, ResolvedConfig } from 'vitest/node'
 import type { ScreenshotOptions } from '../../../context'
-import { mkdir } from 'node:fs/promises'
+import { mkdir, rm } from 'node:fs/promises'
 import { normalize } from 'node:path'
 import { basename, dirname, relative, resolve } from 'pathe'
 import { PlaywrightBrowserProvider } from '../providers/playwright'
@@ -15,6 +15,12 @@ export const screenshot: BrowserCommand<[string, ScreenshotOptions]> = async (
     throw new Error(`Cannot take a screenshot without a test path`)
   }
 
+  options.save ??= true
+
+  if (!options.save) {
+    options.base64 = true
+  }
+
   const path = options.path
     ? resolve(dirname(context.testPath), options.path)
     : resolveScreenshotPath(
@@ -31,28 +37,28 @@ export const screenshot: BrowserCommand<[string, ScreenshotOptions]> = async (
       const element = context.iframe.locator(`${selector}`)
       const buffer = await element.screenshot({
         ...config,
-        path: savePath,
+        path: options.save ? savePath : undefined,
       })
       return returnResult(options, path, buffer)
     }
 
     const buffer = await context.iframe.locator('body').screenshot({
       ...options,
-      path: savePath,
+      path: options.save ? savePath : undefined,
     })
     return returnResult(options, path, buffer)
   }
 
   if (context.provider instanceof WebdriverBrowserProvider) {
     const page = context.provider.browser!
-    if (!options.element) {
-      const body = await page.$('body')
-      const buffer = await body.saveScreenshot(savePath)
-      return returnResult(options, path, buffer)
-    }
+    const element = !options.element
+      ? await page.$('body')
+      : await page.$(`${options.element}`)
 
-    const element = await page.$(`${options.element}`)
     const buffer = await element.saveScreenshot(savePath)
+    if (!options.save) {
+      await rm(savePath, { force: true })
+    }
     return returnResult(options, path, buffer)
   }
 
@@ -84,6 +90,9 @@ function returnResult(
   path: string,
   buffer: Buffer,
 ) {
+  if (!options.save) {
+    return buffer.toString('base64')
+  }
   if (options.base64) {
     return { path, base64: buffer.toString('base64') }
   }

test/browser/test/dom.test.ts

@@ -65,6 +65,20 @@ describe('dom related activity', () => {
     expect(base64).toBeTypeOf('string')
   })
 
+  test('doesn\'t save base64', async () => {
+    const wrapper = createWrapper()
+    const div = createNode()
+    wrapper.appendChild(div)
+
+    const base64 = await page.screenshot({
+      element: wrapper,
+      save: false,
+    })
+    expect(base64).toBeTypeOf('string')
+    expect(base64).not.toContain('__screenshots__')
+    expect(base64).not.toContain('dom.test.ts')
+  })
+
   test('shadow dom screenshot', async () => {
     const wrapper = createWrapper()
     const div = createNode()