feat: add an option to return base64 from page.screenshot (#5993)

sheremet-va · web-flow · commit be32317638c0 · 2024-06-30T19:12:32.000+02:00
diff --git a/docs/guide/browser.md b/docs/guide/browser.md
@@ -465,12 +465,16 @@ export const page: {
   /**
    * Change the size of iframe's viewport.
    */
-  viewport: (width: number | string, height: number | string) => Promise<void>
+  viewport(width: number | string, height: number | string): Promise<void>
   /**
    * Make a screenshot of the test iframe or a specific element.
-   * @returns Path to the screenshot file.
+   * @returns Path to the screenshot file or path and base64.
    */
-  screenshot: (options?: ScreenshotOptions) => Promise<string>
+  screenshot(options: Omit<ScreenshotOptions, 'base64'> & { base64: true }): Promise<{
+    path: string
+    base64: string
+  }>
+  screenshot(options?: ScreenshotOptions): Promise<string>
 }
 
 export const cdp: () => CDPSession
diff --git a/eslint.config.js b/eslint.config.js
@@ -100,6 +100,7 @@ export default antfu(
       'import/newline-after-import': 'off',
       'import/first': 'off',
       'unused-imports/no-unused-imports': 'off',
+      'ts/method-signature-style': 'off',
     },
   },
   {
diff --git a/packages/browser/context.d.ts b/packages/browser/context.d.ts
@@ -29,6 +29,10 @@ export interface ScreenshotOptions {
    * Path relative to the `screenshotDirectory` in the test config.
    */
   path?: string
+  /**
+   * Will also return the base64 encoded screenshot alongside the path.
+   */
+  base64?: boolean
 }
 
 export interface BrowserCommands {
@@ -245,12 +249,16 @@ export interface BrowserPage {
   /**
    * Change the size of iframe's viewport.
    */
-  viewport: (width: number, height: number) => Promise<void>
+  viewport(width: number, height: number): Promise<void>
   /**
    * Make a screenshot of the test iframe or a specific element.
-   * @returns Path to the screenshot file.
+   * @returns Path to the screenshot file or path and base64.
    */
-  screenshot: (options?: ScreenshotOptions) => Promise<string>
+  screenshot(options: Omit<ScreenshotOptions, 'base64'> & { base64: true }): Promise<{
+    path: string
+    base64: string
+  }>
+  screenshot(options?: ScreenshotOptions): Promise<string>
 }
 
 export const page: BrowserPage
diff --git a/packages/browser/src/node/commands/screenshot.ts b/packages/browser/src/node/commands/screenshot.ts
@@ -29,28 +29,29 @@ export const screenshot: BrowserCommand<[string, ScreenshotOptions]> = async (
     if (options.element) {
       const { element: elementXpath, ...config } = options
       const element = context.iframe.locator(`xpath=${elementXpath}`)
-      await element.screenshot({ ...config, path: savePath })
+      const buffer = await element.screenshot({ ...config, path: savePath })
+      return returnResult(options, path, buffer)
     }
-    else {
-      await context.iframe.locator('body').screenshot({
-        ...options,
-        path: savePath,
-      })
-    }
-    return path
+
+    const buffer = await context.iframe.locator('body').screenshot({
+      ...options,
+      path: savePath,
+    })
+    return returnResult(options, path, buffer)
   }
 
   if (context.provider instanceof WebdriverBrowserProvider) {
     const page = context.provider.browser!
     if (!options.element) {
       const body = await page.$('body')
-      await body.saveScreenshot(savePath)
-      return path
+      const buffer = await body.saveScreenshot(savePath)
+      return returnResult(options, path, buffer)
     }
+
     const xpath = `//${options.element}`
     const element = await page.$(xpath)
-    await element.saveScreenshot(savePath)
-    return path
+    const buffer = await element.saveScreenshot(savePath)
+    return returnResult(options, path, buffer)
   }
 
   throw new Error(
@@ -75,3 +76,14 @@ function resolveScreenshotPath(
   }
   return resolve(dir, '__screenshots__', base, name)
 }
+
+function returnResult(
+  options: ScreenshotOptions,
+  path: string,
+  buffer: Buffer,
+) {
+  if (options.base64) {
+    return { path, base64: buffer.toString('base64') }
+  }
+  return path
+}
diff --git a/test/browser/test/dom.test.ts b/test/browser/test/dom.test.ts
@@ -1,16 +1,19 @@
-import { describe, expect, test } from 'vitest'
+import { beforeEach, describe, expect, test } from 'vitest'
 import { page } from '@vitest/browser/context'
 import { createNode } from '#src/createNode'
 import '../src/button.css'
 
 describe('dom related activity', () => {
-  test('renders div', async () => {
+  beforeEach(() => {
     document.body.style.background = '#f3f3f3'
-    const wrapper = document.createElement('div')
-    wrapper.className = 'wrapper'
-    document.body.appendChild(wrapper)
+    document.body.replaceChildren()
+  })
+
+  test('renders div', async () => {
+    const wrapper = createWrapper()
     const div = createNode()
     wrapper.appendChild(div)
+
     await expect.element(div).toHaveTextContent('Hello World!')
     const screenshotPath = await page.screenshot({
       element: wrapper,
@@ -19,4 +22,27 @@ describe('dom related activity', () => {
       /__screenshots__\/dom.test.ts\/dom-related-activity-renders-div-1.png/,
     )
   })
+
+  test('resolves base64 screenshot', async () => {
+    const wrapper = createWrapper()
+    const div = createNode()
+    wrapper.appendChild(div)
+
+    const { path, base64 } = await page.screenshot({
+      element: wrapper,
+      base64: true,
+    })
+    expect(path).toMatch(
+      /__screenshots__\/dom.test.ts\/dom-related-activity-resolves-base64-screenshot-1.png/,
+    )
+    expect(base64).toBeTypeOf('string')
+  })
 })
+
+function createWrapper() {
+  document.body.style.background = '#f3f3f3'
+  const wrapper = document.createElement('div')
+  wrapper.className = 'wrapper'
+  document.body.appendChild(wrapper)
+  return wrapper
+}

Original file line number	Diff line number	Diff line change
`@@ -100,6 +100,7 @@ export default antfu(`
`100`	`100`	`'import/newline-after-import': 'off',`
`101`	`101`	`'import/first': 'off',`
`102`	`102`	`'unused-imports/no-unused-imports': 'off',`
	`103`	`+ 'ts/method-signature-style': 'off',`
`103`	`104`	`},`
`104`	`105`	`},`
`105`	`106`	`{`