From d118088fa7284be516b3f9871036e108ed175801 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Greg=20Berg=C3=A9?= <berge.greg@gmail.com>
Date: Wed, 10 Jan 2024 16:07:16 +0100
Subject: [PATCH] feat(puppeteer): allow to specify the full path of screenshot
 (#103)

---
 packages/puppeteer/docs/index.mdx |  2 +-
 packages/puppeteer/src/index.ts   | 27 ++++++++++++++++++++-------
 packages/webdriverio/src/index.ts | 19 +++++++++++++------
 3 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/packages/puppeteer/docs/index.mdx b/packages/puppeteer/docs/index.mdx
index 90be1f7e..a20f5ad6 100644
--- a/packages/puppeteer/docs/index.mdx
+++ b/packages/puppeteer/docs/index.mdx
@@ -51,7 +51,7 @@ Screenshots are stored in `screenshots/argos` folder, relative to current direct
 ### argosScreenshot(page, name[, options])
 
 - `page` - A `puppeteer` page instance
-- `name` - The screenshot name; must be unique
+- `name` - The screenshot name; must be unique. If ends by `.png` we treat it as a path.
 - `options` - See [Page.screenshot command options](https://pptr.dev/next/api/puppeteer.page.screenshot/)
 - `options.element` - Accept an ElementHandle or a string selector to screenshot an element
 - `options.viewports` - Specifies the viewports for which to capture screenshots. See [viewports configuration](/viewports).
diff --git a/packages/puppeteer/src/index.ts b/packages/puppeteer/src/index.ts
index 8fcfb5f6..27a1d16d 100644
--- a/packages/puppeteer/src/index.ts
+++ b/packages/puppeteer/src/index.ts
@@ -68,6 +68,22 @@ async function getBrowserInfo(page: Page) {
   return { browserName, browserVersion };
 }
 
+async function getScreenshotPath(name: string) {
+  if (name.endsWith(".png")) return name;
+
+  const screenshotFolder = resolve(process.cwd(), "screenshots/argos");
+  await mkdir(screenshotFolder, { recursive: true });
+  return resolve(screenshotFolder, name + ".png");
+}
+
+/**
+ * @param page Puppeteer `page` object.
+ * @param name The name of the screenshot or the full path to the screenshot.
+ * @param options In addition to Puppeteer's `ScreenshotOptions`, you can pass:
+ * @param options.element ElementHandle or string selector of the element to take a screenshot of.
+ * @param options.viewports Viewports to take screenshots of.
+ * @param options.argosCSS Custom CSS evaluated during the screenshot process.
+ */
 export async function argosScreenshot(
   page: Page,
   name: string,
@@ -80,12 +96,8 @@ export async function argosScreenshot(
     throw new Error("The `name` argument is required.");
   }
 
-  const screenshotFolder = resolve(process.cwd(), "screenshots/argos");
-
   const [originalViewport] = await Promise.all([
     getViewport(page),
-    // Create the screenshot folder if it doesn't exist
-    mkdir(screenshotFolder, { recursive: true }),
     // Inject Argos script into the page
     injectArgos(page),
   ]);
@@ -148,9 +160,10 @@ export async function argosScreenshot(
       ((window as any).__ARGOS__ as ArgosGlobal).waitForStability(),
     );
 
-    const screenshotPath = resolve(screenshotFolder, `${name}.png`);
-
-    const metadata = await collectMetadata();
+    const [screenshotPath, metadata] = await Promise.all([
+      getScreenshotPath(name),
+      collectMetadata(),
+    ]);
 
     await writeMetadata(screenshotPath, metadata);
 
diff --git a/packages/webdriverio/src/index.ts b/packages/webdriverio/src/index.ts
index 89f81279..9c144999 100644
--- a/packages/webdriverio/src/index.ts
+++ b/packages/webdriverio/src/index.ts
@@ -67,7 +67,7 @@ async function getImageDimensions(buffer: Buffer): Promise<Size> {
   }
 }
 
-async function getFilePath(name: string) {
+async function getScreenshotPath(name: string) {
   if (name.endsWith(".png")) return name;
 
   const screenshotFolder = resolve(process.cwd(), "screenshots/argos");
@@ -93,8 +93,15 @@ export type ArgosScreenshotOptions = {
 /**
  * Take a screenshot of the current page, optionally masking certain areas.
  * @param browser A WebdriverIO `browser` object.
- * @param filepath The path to save the screenshot to.
- * @param options Options for the screenshot.
+ * @param name The name of the screenshot or the full path to the screenshot.
+ * @param options Options to customize the screenshot.
+ * @param options.mask
+ * Specify ares that should be masked when the screenshot is taken.
+ * Masked elements will be overlaid with a pink box #FF00FF (customized by maskColor)
+ * that completely covers its bounding box.
+ * @param options.maskColor
+ * Specify the color of the overlay box for masked elements, in CSS color format.
+ * Default color is pink #FF00FF.
  */
 export async function argosScreenshot(
   browser: WebdriverIO.Browser,
@@ -108,8 +115,8 @@ export async function argosScreenshot(
     throw new Error("The `name` argument is required.");
   }
 
-  const filepath = await getFilePath(name);
-  const imageBuffer = await browser.saveScreenshot(filepath);
+  const screenshotPath = await getScreenshotPath(name);
+  const imageBuffer = await browser.saveScreenshot(screenshotPath);
 
   if (!mask) {
     return imageBuffer;
@@ -118,6 +125,6 @@ export async function argosScreenshot(
   const dimensions = await getImageDimensions(imageBuffer);
   const maskBuffer = await createMask(dimensions, mask, maskColor);
   const maskedBuffer = await applyMask(imageBuffer, maskBuffer);
-  await sharp(maskedBuffer).toFile(filepath);
+  await sharp(maskedBuffer).toFile(screenshotPath);
   return maskedBuffer;
 }