feat(webdriverio): add hide, remove and selector masks

packages/webdriverio/fixtures/helpers.html

@@ -0,0 +1,37 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8" />
+    <title>Helpers page</title>
+    <meta name="viewport" content="width=device-width,initial-scale=1" />
+    <style>
+      body {
+        margin: 0;
+        background: white;
+      }
+      div {
+        height: 50px;
+      }
+      footer {
+        height: 2000px;
+      }
+    </style>
+  </head>
+  <body>
+    <div id="hidden" style="background: red"></div>
+    <div id="removed" style="background: blue"></div>
+    <div
+      id="transparent"
+      data-visual-test="transparent"
+      style="background: lime"
+    ></div>
+    <div
+      id="blackout"
+      data-visual-test="blackout"
+      style="background: yellow"
+    ></div>
+    <div id="masked" style="background: cyan"></div>
+    <div id="last" style="background: black"></div>
+    <footer></footer>
+  </body>
+</html>

packages/webdriverio/package.json

@@ -48,6 +48,7 @@
     "webdriverio": "^6 || ^7 || ^8 || ^9 || ^10"
   },
   "dependencies": {
+    "@argos-ci/browser": "workspace:*",
     "sharp": "^0.35.5"
   },
   "devDependencies": {

packages/webdriverio/src/index.ts

@@ -1,72 +1,71 @@
 import sharp from "sharp";
 import { resolve } from "node:path";
-import { mkdir } from "node:fs/promises";
+import { mkdir, writeFile } from "node:fs/promises";
+import {
+  type ArgosGlobal,
+  getGlobalScript,
+  type StabilizationContext,
+} from "@argos-ci/browser";
 
 type Size = { width: number; height: number };
 type Coords = { x: number; y: number; width: number; height: number };
 
-// Function to create a mask with Sharp
-async function createMask(
-  dimensions: Size,
+/**
+ * Elements with this attribute are always masked.
+ */
+const BLACKOUT_SELECTOR = '[data-visual-test="blackout"]';
+
+/**
+ * Clip an area to the image, on whole pixels.
+ * Returns `null` when the area is outside of the image.
+ */
+function clipArea(area: Coords, dimensions: Size): Coords | null {
+  const left = Math.max(0, Math.floor(area.x));
+  const top = Math.max(0, Math.floor(area.y));
+  const right = Math.min(dimensions.width, Math.ceil(area.x + area.width));
+  const bottom = Math.min(dimensions.height, Math.ceil(area.y + area.height));
+  if (right <= left || bottom <= top) {
+    return null;
+  }
+  return { x: left, y: top, width: right - left, height: bottom - top };
+}
+
+/**
+ * Overlay areas of the screenshot with the mask color.
+ */
+async function maskScreenshot(
+  image: Buffer,
   areas: Coords[],
   maskColor: string,
 ) {
-  // Start with a black image
-  let mask = sharp({
-    create: {
-      width: dimensions.width,
-      height: dimensions.height,
-      channels: 4,
-      background: { r: 0, g: 0, b: 0, alpha: 0 },
-    },
-  });
-
-  const areasImages = areas.map((coords) => ({
-    input: Buffer.from(
-      '<svg><rect x="0" y="0" width="' +
-        coords.width +
-        '" height="' +
-        coords.height +
-        `" fill="${maskColor}"/></svg>`,
-    ),
-    top: coords.y,
-    left: coords.x,
-  }));
-
-  mask = mask.composite(areasImages);
-
-  return mask.png().toBuffer();
-}
-
-async function applyMask(image: Buffer, mask: Buffer) {
   try {
-    return await sharp(image)
-      .composite([
-        {
-          input: mask,
-          blend: "over",
-        },
-      ])
+    const screenshot = sharp(image);
+    const dimensions = await screenshot.metadata();
+    return await screenshot
+      .composite(
+        areas
+          // Sharp can't composite an image larger than the screenshot.
+          .map((area) => clipArea(area, dimensions))
+          .filter((coords) => coords !== null)
+          .map((coords) => ({
+            input: Buffer.from(
+              '<svg><rect x="0" y="0" width="' +
+                coords.width +
+                '" height="' +
+                coords.height +
+                `" fill="${maskColor}"/></svg>`,
+            ),
+            top: coords.y,
+            left: coords.x,
+          })),
+      )
       .png()
       .toBuffer();
   } catch (err) {
     throw new Error("Error applying mask", { cause: err });
   }
 }
 
-async function getImageDimensions(buffer: Buffer): Promise<Size> {
-  try {
-    const imageInfo = await sharp(buffer).metadata();
-    const { width, height } = imageInfo;
-    if (!width || !height) {
-      throw new Error("Dimensions not found.");
-    }
-    return { width, height };
-  } catch (err) {
-    throw new Error("Error getting image dimensions", { cause: err });
-  }
-}
-
 async function getScreenshotPath(name: string) {
   if (name.endsWith(".png")) {
     return name;
@@ -79,54 +78,203 @@ async function getScreenshotPath(name: string) {
 
 export type ArgosScreenshotOptions = {
   /**
-   * 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.
+   * Areas to mask when the screenshot is taken, as rectangles in pixels of the
+   * screenshot or as WebdriverIO selectors of the elements to mask.
+   * Masked areas are overlaid with a pink box #FF00FF (customized by `maskColor`).
+   * Elements with the `data-visual-test="blackout"` attribute are always masked.
+   * Selectors that match nothing are ignored.
+   * @example
+   *    argosScreenshot(browser, "home", {
+   *      mask: ["#clock", { x: 0, y: 0, width: 100, height: 100 }],
+   *    })
    */
-  mask?: { x: number; y: number; width: number; height: number }[];
+  mask?: (string | Coords)[];
+
   /**
    * Specify the color of the overlay box for masked elements, in CSS color format.
    * Default color is pink #FF00FF.
    * @default "#FF00FF"
    */
   maskColor?: string;
+
+  /**
+   * WebdriverIO selectors of the elements to hide from the screenshot.
+   * Hidden elements keep their space in the layout, like with the
+   * `data-visual-test="transparent"` helper.
+   * Selectors that match nothing are ignored.
+   * @example
+   *    argosScreenshot(browser, "home", { hide: ["#clock", ".avatar"] })
+   */
+  hide?: string[];
+
+  /**
+   * WebdriverIO selectors of the elements to remove from the screenshot.
+   * Removed elements no longer take space in the layout, like with the
+   * `data-visual-test="removed"` helper.
+   * Selectors that match nothing are ignored.
+   * @example
+   *    argosScreenshot(browser, "home", { remove: [".cookie-banner"] })
+   */
+  remove?: string[];
+
+  /**
+   * Custom CSS evaluated during the screenshot process.
+   */
+  argosCSS?: string;
 };
 
 /**
- * Take a screenshot of the current page, optionally masking certain areas.
- * @param browser A WebdriverIO `browser` object.
- * @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.
+ * Find the elements matched by WebdriverIO selectors.
+ */
+async function findElements(
+  browser: WebdriverIO.Browser,
+  selectors: string[] = [],
+) {
+  const results = await Promise.all(
+    selectors.map((selector) => browser.$$(selector)),
+  );
+  return results.flatMap((elements) => [...elements]);
+}
+
+/**
+ * Inject Argos script into the page.
+ * Checked in the page, to save a round trip to the browser.
+ */
+async function injectArgos(browser: WebdriverIO.Browser) {
+  await browser.execute(
+    `if (typeof window.__ARGOS__ === "undefined") {\n${getGlobalScript()}\n}`,
+  );
+}
+
+/**
+ * Prepare the page for the screenshot: add the `__argos__` class, inject
+ * `argosCSS`, enable the `data-visual-test` helpers and apply the `hide` and
+ * `remove` options. The stabilization plugins are not run.
+ * Returns a function restoring the page.
+ */
+async function setup(
+  browser: WebdriverIO.Browser,
+  options: ArgosScreenshotOptions,
+) {
+  const context: StabilizationContext = {
+    argosCSS: options.argosCSS,
+    options: false,
+  };
+
+  const [, hidden, removed] = await Promise.all([
+    injectArgos(browser),
+    findElements(browser, options.hide),
+    findElements(browser, options.remove),
+  ]);
+
+  await browser.execute(
+    (context, hiddenCount, ...elements) => {
+      const argos = (window as any).__ARGOS__ as ArgosGlobal;
+      argos.beforeAll(context);
+      argos.applyHelper(elements.slice(0, hiddenCount), "transparent");
+      argos.applyHelper(elements.slice(hiddenCount), "removed");
+      argos.beforeEach(context);
+    },
+    context,
+    hidden.length,
+    ...hidden,
+    ...removed,
+  );
+
+  return async () => {
+    await browser.execute(() => {
+      const argos = (window as any).__ARGOS__ as ArgosGlobal;
+      argos.afterEach();
+      argos.afterAll();
+    });
+  };
+}
+
+/**
+ * Get the areas covered by the elements to mask, in pixels of the screenshot.
+ * The screenshot only captures the viewport.
+ */
+async function getMaskedElementsAreas(
+  browser: WebdriverIO.Browser,
+  selectors: string[],
+): Promise<Coords[]> {
+  const elements = await findElements(browser, selectors);
+  if (elements.length === 0) {
+    return [];
+  }
+  return browser.execute(
+    (...elements) =>
+      elements.map((element) => {
+        const rect = element.getBoundingClientRect();
+        const ratio = window.devicePixelRatio;
+        return {
+          x: rect.x * ratio,
+          y: rect.y * ratio,
+          width: rect.width * ratio,
+          height: rect.height * ratio,
+        };
+      }),
+    ...elements,
+  );
+}
+
+/**
+ * Take a screenshot of the current page, optionally hiding, removing or
+ * masking some elements.
+ * The `data-visual-test` helper attributes are applied.
+ *
+ * @example
+ *    argosScreenshot(browser, "my-screenshot")
+ * @see https://argos-ci.com/docs/reference/webdriverio
  */
 export async function argosScreenshot(
+  /**
+   * WebdriverIO `browser` object.
+   */
   browser: WebdriverIO.Browser,
+  /**
+   * Name of the screenshot or full path to the screenshot.
+   */
   name: string,
-  { mask, maskColor = "#FF00FF" }: ArgosScreenshotOptions = {},
+  /**
+   * Options for the screenshot.
+   */
+  options: ArgosScreenshotOptions = {},
 ) {
+  const { mask = [], maskColor = "#FF00FF" } = options;
+
   if (!browser) {
     throw new Error("A WebdriverIO `browser` object is required.");
   }
   if (!name) {
     throw new Error("The `name` argument is required.");
   }
 
+  const teardown = await setup(browser, options);
+
+  // Measured once the elements are hidden or removed, as the layout may change.
+  const maskedElementsAreas = await getMaskedElementsAreas(browser, [
+    BLACKOUT_SELECTOR,
+    ...mask.filter((area) => typeof area === "string"),
+  ]);
+  const areas = [
+    ...mask.filter((area) => typeof area !== "string"),
+    ...maskedElementsAreas,
+  ];
+
   const screenshotPath = await getScreenshotPath(name);
   const imageBuffer = await browser.saveScreenshot(screenshotPath);
 
-  if (!mask) {
+  if (areas.length === 0) {
+    await teardown();
     return imageBuffer;
   }
 
-  const dimensions = await getImageDimensions(imageBuffer);
-  const maskBuffer = await createMask(dimensions, mask, maskColor);
-  const maskedBuffer = await applyMask(imageBuffer, maskBuffer);
-  await sharp(maskedBuffer).toFile(screenshotPath);
+  // Restore the page while masking the screenshot.
+  const [maskedBuffer] = await Promise.all([
+    maskScreenshot(imageBuffer, areas, maskColor),
+    teardown(),
+  ]);
+  await writeFile(screenshotPath, maskedBuffer);
   return maskedBuffer;
 }