feat: reset argos changes after screenshot

packages/browser/src/global.ts

@@ -1,14 +1,17 @@
 import {
   waitForStability,
-  prepareForScreenshot,
+  setup,
+  teardown,
   PrepareForScreenshotOptions,
 } from "./stabilization";
 import { getColorScheme, getMediaType } from "./media";
 
 const ArgosGlobal = {
   waitForStability: () => waitForStability(document),
-  prepareForScreenshot: (options: PrepareForScreenshotOptions = {}) =>
-    prepareForScreenshot(document, options),
+  setup: (options: PrepareForScreenshotOptions = {}) =>
+    setup(document, options),
+  teardown: (options: PrepareForScreenshotOptions = {}) =>
+    teardown(document, options),
   getColorScheme: () => getColorScheme(window),
   getMediaType: () => getMediaType(window),
 };

packages/browser/src/stabilization.ts

@@ -34,14 +34,44 @@ export const GLOBAL_CSS: string = `
 }
 `;
 
+/**
+ * Set the position of an element and backup the previous one.
+ */
+function setAndBackupSpellcheck(element: HTMLElement, spellcheck: string) {
+  element.setAttribute(
+    "data-argos-bck-position",
+    element.getAttribute("spellcheck") ?? "unset",
+  );
+  element.setAttribute("spellcheck", spellcheck);
+}
+
+const SPELL_CHECK_QUERY =
+  "[contenteditable]:not([contenteditable=false]):not([spellcheck=false]), input:not([spellcheck=false]), textarea:not([spellcheck=false])";
+
 /**
  * Disable spellcheck to avoid displaying markers.
  */
 export function disableSpellCheck(document: Document) {
-  const query =
-    "[contenteditable]:not([contenteditable=false]):not([spellcheck=false]), input:not([spellcheck=false]), textarea:not([spellcheck=false])";
-  const inputs = document.querySelectorAll(query);
-  inputs.forEach((input) => input.setAttribute("spellcheck", "false"));
+  const inputs = document.querySelectorAll(SPELL_CHECK_QUERY);
+  inputs.forEach((element) => {
+    if (!checkIsHTMLElement(element)) return;
+    setAndBackupSpellcheck(element, "false");
+  });
+}
+
+/**
+ * Restore spellcheck attribute.
+ */
+export function restoreSpellCheck(document: Document) {
+  const inputs = document.querySelectorAll(SPELL_CHECK_QUERY);
+  inputs.forEach((input) => {
+    const spellcheck = input.getAttribute("data-argos-bck-spellcheck");
+    if (spellcheck === "unset") {
+      input.removeAttribute("spellcheck");
+    } else if (spellcheck) {
+      input.setAttribute("spellcheck", spellcheck);
+    }
+  });
 }
 
 /**
@@ -50,13 +80,35 @@ export function disableSpellCheck(document: Document) {
 export function injectGlobalStyles(document: Document) {
   const style = document.createElement("style");
   style.textContent = GLOBAL_CSS;
+  style.id = "argos-global-styles";
   document.head.appendChild(style);
 }
 
+/**
+ * Remove global styles from the DOM.
+ */
+export function removeGlobalStyles(document: Document) {
+  const style = document.getElementById("argos-global-styles");
+  if (style) {
+    style.remove();
+  }
+}
+
 const checkIsHTMLElement = (element: Element): element is HTMLElement => {
   return "style" in element;
 };
 
+/**
+ * Set the position of an element and backup the previous one.
+ */
+function setAndBackupPosition(element: HTMLElement, position: string) {
+  element.setAttribute(
+    "data-argos-bck-position",
+    element.style.position ?? "unset",
+  );
+  element.style.position = position;
+}
+
 /**
  * Stabilize sticky and fixed elements.
  */
@@ -69,19 +121,37 @@ export function stabilizeElementPositions(document: Document) {
     const style = window.getComputedStyle(element);
     const position = style.position;
     if (position === "fixed") {
-      element.style.position = "absolute";
+      setAndBackupPosition(element, "absolute");
     } else if (position === "sticky") {
-      element.style.position = "relative";
+      setAndBackupPosition(element, "relative");
+    }
+  });
+}
+
+/**
+ * Restore the position of elements.
+ */
+export function restoreElementPositions(document: Document) {
+  const window = document.defaultView;
+  if (!window) return;
+  const elements = Array.from(document.querySelectorAll("*"));
+  elements.forEach((element) => {
+    if (!checkIsHTMLElement(element)) return;
+    const position = element.getAttribute("data-argos-bck-position");
+    if (position === "unset") {
+      element.style.removeProperty("position");
+    } else if (position) {
+      element.style.position = position;
     }
   });
 }
 
 export type PrepareForScreenshotOptions = { fullPage?: boolean };
 
 /**
- * Prepare the document for a screenshot.
+ * Setup the document for screenshots.
  */
-export function prepareForScreenshot(
+export function setup(
   document: Document,
   { fullPage }: PrepareForScreenshotOptions = {},
 ) {
@@ -92,6 +162,20 @@ export function prepareForScreenshot(
   }
 }
 
+/**
+ * Restore the document after screenshots.
+ */
+export function teardown(
+  document: Document,
+  { fullPage }: PrepareForScreenshotOptions = {},
+) {
+  removeGlobalStyles(document);
+  restoreSpellCheck(document);
+  if (fullPage) {
+    restoreElementPositions(document);
+  }
+}
+
 /**
  * Wait for the fonts to be loaded.
  */

packages/cypress/src/support.ts

@@ -70,9 +70,7 @@ Cypress.Commands.add(
     const fullPage = !options.capture || options.capture === "fullPage";
 
     cy.window({ log: false }).then((window) =>
-      ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot({
-        fullPage,
-      }),
+      ((window as any).__ARGOS__ as ArgosGlobal).setup({ fullPage }),
     );
 
     function stabilizeAndScreenshot(name: string) {
@@ -152,6 +150,11 @@ Cypress.Commands.add(
       );
     }
 
+    // Teardown Argos
+    cy.window({ log: false }).then((window) =>
+      ((window as any).__ARGOS__ as ArgosGlobal).teardown({ fullPage }),
+    );
+
     // Restore the original viewport
     cy.viewport(
       Cypress.config("viewportWidth"),

packages/playwright/src/index.ts

@@ -105,9 +105,7 @@ export async function argosScreenshot(
 
   await page.evaluate(
     ({ fullPage }) =>
-      ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot({
-        fullPage,
-      }),
+      ((window as any).__ARGOS__ as ArgosGlobal).setup({ fullPage }),
     { fullPage },
   );
 
@@ -211,6 +209,13 @@ export async function argosScreenshot(
     );
   }
 
+  // Teardown Argos
+  await page.evaluate(
+    ({ fullPage }) =>
+      ((window as any).__ARGOS__ as ArgosGlobal).teardown({ fullPage }),
+    { fullPage },
+  );
+
   // Restore the original viewport size
   await page.setViewportSize(originalViewportSize);
 }

packages/puppeteer/src/index.ts

@@ -87,9 +87,7 @@ export async function argosScreenshot(
 
   await page.evaluate(
     ({ fullPage }) =>
-      ((window as any).__ARGOS__ as ArgosGlobal).prepareForScreenshot({
-        fullPage,
-      }),
+      ((window as any).__ARGOS__ as ArgosGlobal).setup({ fullPage }),
     { fullPage },
   );
 
@@ -191,6 +189,13 @@ export async function argosScreenshot(
     );
   }
 
+  // Teardown Argos
+  await page.evaluate(
+    ({ fullPage }) =>
+      ((window as any).__ARGOS__ as ArgosGlobal).teardown({ fullPage }),
+    { fullPage },
+  );
+
   // Restore the original viewport
   await page.setViewport(originalViewport);
 }