From 5f01dd962a3a7a010eb2df8340d37e9d720c250b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Greg=20Berg=C3=A9?= <berge.greg@gmail.com>
Date: Thu, 16 Nov 2023 18:14:25 +0100
Subject: [PATCH] feat: stabilize sticky & fixed elements in full page

---
 packages/browser/src/global.ts        |  9 +++++--
 packages/browser/src/stabilization.ts | 35 +++++++++++++++++++++++++--
 packages/cypress/src/support.ts       |  6 ++++-
 packages/playwright/src/index.ts      |  9 +++++--
 packages/puppeteer/src/index.ts       |  9 +++++--
 5 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/packages/browser/src/global.ts b/packages/browser/src/global.ts
index d17ffe4..25e006a 100644
--- a/packages/browser/src/global.ts
+++ b/packages/browser/src/global.ts
@@ -1,9 +1,14 @@
-import { waitForStability, prepareForScreenshot } from "./stabilization";
+import {
+  waitForStability,
+  prepareForScreenshot,
+  PrepareForScreenshotOptions,
+} from "./stabilization";
 import { getColorScheme, getMediaType } from "./media";
 
 const ArgosGlobal = {
   waitForStability: () => waitForStability(document),
-  prepareForScreenshot: () => prepareForScreenshot(document),
+  prepareForScreenshot: (options: PrepareForScreenshotOptions = {}) =>
+    prepareForScreenshot(document, options),
   getColorScheme: () => getColorScheme(window),
   getMediaType: () => getMediaType(window),
 };
diff --git a/packages/browser/src/stabilization.ts b/packages/browser/src/stabilization.ts
index d09712b..3298ffb 100644
--- a/packages/browser/src/stabilization.ts
+++ b/packages/browser/src/stabilization.ts
@@ -45,7 +45,7 @@ export function disableSpellCheck(document: Document) {
 }
 
 /**
- *
+ * Inject global styles in the DOM.
  */
 export function injectGlobalStyles(document: Document) {
   const style = document.createElement("style");
@@ -53,12 +53,43 @@ export function injectGlobalStyles(document: Document) {
   document.head.appendChild(style);
 }
 
+const checkIsHTMLElement = (element: Element): element is HTMLElement => {
+  return "style" in element;
+};
+
+/**
+ * Stabilize sticky and fixed elements.
+ */
+export function stabilizeElementPositions(document: Document) {
+  const window = document.defaultView;
+  if (!window) return;
+  const elements = Array.from(document.querySelectorAll("*"));
+  elements.forEach((element) => {
+    if (!checkIsHTMLElement(element)) return;
+    const style = window.getComputedStyle(element);
+    const position = style.position;
+    if (position === "fixed") {
+      element.style.position = "absolute";
+    } else if (position === "sticky") {
+      element.style.position = "relative";
+    }
+  });
+}
+
+export type PrepareForScreenshotOptions = { fullPage?: boolean };
+
 /**
  * Prepare the document for a screenshot.
  */
-export function prepareForScreenshot(document: Document) {
+export function prepareForScreenshot(
+  document: Document,
+  { fullPage }: PrepareForScreenshotOptions = {},
+) {
   injectGlobalStyles(document);
   disableSpellCheck(document);
+  if (fullPage) {
+    stabilizeElementPositions(document);
+  }
 }
 
 /**
diff --git a/packages/cypress/src/support.ts b/packages/cypress/src/support.ts
index 34e5ee7..58347b2 100644
--- a/packages/cypress/src/support.ts
+++ b/packages/cypress/src/support.ts
@@ -67,8 +67,12 @@ Cypress.Commands.add(
 
     injectArgos();
 
+    const fullPage = !options.capture || options.capture === "fullPage";
+
     cy.window({ log: false }).then((window) =>
-      ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot(),
+      ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot({
+        fullPage,
+      }),
     );
 
     function stabilizeAndScreenshot(name: string) {
diff --git a/packages/playwright/src/index.ts b/packages/playwright/src/index.ts
index 336fc58..2464a79 100644
--- a/packages/playwright/src/index.ts
+++ b/packages/playwright/src/index.ts
@@ -100,8 +100,13 @@ export async function argosScreenshot(
 
   const originalViewportSize = getViewportSize(page);
 
+  const fullPage =
+    options.fullPage !== undefined ? options.fullPage : handle === page;
+
   await page.evaluate(() =>
-    ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot(),
+    ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot({
+      fullPage,
+    }),
   );
 
   async function collectMetadata(
@@ -165,7 +170,7 @@ export async function argosScreenshot(
       handle.screenshot({
         path: screenshotPath ?? undefined,
         type: "png",
-        fullPage: handle === page,
+        fullPage,
         mask: [page.locator('[data-visual-test="blackout"]')],
         animations: "disabled",
         ...options,
diff --git a/packages/puppeteer/src/index.ts b/packages/puppeteer/src/index.ts
index 833935e..a0c31e0 100644
--- a/packages/puppeteer/src/index.ts
+++ b/packages/puppeteer/src/index.ts
@@ -83,9 +83,14 @@ export async function argosScreenshot(
   ]);
 
   await page.evaluate(() =>
-    ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot(),
+    ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot({
+      fullPage,
+    }),
   );
 
+  const fullPage =
+    options.fullPage !== undefined ? options.fullPage : element === undefined;
+
   async function collectMetadata(): Promise<ScreenshotMetadata> {
     const [
       colorScheme,
@@ -144,7 +149,7 @@ export async function argosScreenshot(
     const screenshotOptions: ScreenshotOptions = {
       path: screenshotPath,
       type: "png",
-      fullPage: element === undefined,
+      fullPage,
       ...options,
     };