From 84382bdc68f3528c916a70edbf086b384f1e7b44 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Mon, 15 Jun 2026 13:16:09 +0200
Subject: [PATCH 01/23] Optimize when no need to crop

---
 src/image-utils/compare.ts     |  85 +++++++++++++++++--------
 src/page-objects/screenshot.ts | 110 +++++++++++++++++++++++----------
 src/page-objects/types.ts      |  16 ++++-
 test/page-object.test.ts       |   9 +--
 4 files changed, 150 insertions(+), 70 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index b9655a0..b5be96e 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -2,8 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 import { PNG } from 'pngjs';
 import pixelmatch from 'pixelmatch';
-import { packPng, cropImage } from './utils';
-import { ElementRect, ElementSize, ScreenshotWithOffset } from '../page-objects/types';
+import { packPng, cropImage, parsePng } from './utils';
+import { ElementRect, ElementSize, Screenshot, ScreenshotWithOffset } from '../page-objects/types';
 
 export function compareImages(firstImage: PNG, secondImage: PNG, { width, height }: ElementSize) {
   // This prevents an error thrown from pixelmatch when comparing 0-sized images.
@@ -19,7 +19,7 @@ export function compareImages(firstImage: PNG, secondImage: PNG, { width, height
   return { diffPixels, diffImage };
 }
 
-function normalizeSize(firstScreenshot: ScreenshotWithOffset, secondScreenshot: ScreenshotWithOffset) {
+function normalizeSize(firstScreenshot: Screenshot, secondScreenshot: Screenshot) {
   return {
     height: Math.round(Math.max(firstScreenshot.height, secondScreenshot.height)),
     width: Math.round(Math.max(firstScreenshot.width, secondScreenshot.width)),
@@ -33,6 +33,10 @@ function scaleSize(size: ElementSize, pixelRatio: number) {
   };
 }
 
+function isZeroOffset(screenshot: ScreenshotWithOffset): boolean {
+  return !screenshot.offset;
+}
+
 export interface CropAndCompareResult {
   firstImage: Buffer;
   secondImage: Buffer;
@@ -42,35 +46,53 @@ export interface CropAndCompareResult {
 }
 
 export async function cropAndCompare(
-  firstScreenshot: ScreenshotWithOffset,
-  secondScreenshot: ScreenshotWithOffset
+  firstScreenshot: Screenshot,
+  secondScreenshot: Screenshot
 ): Promise<CropAndCompareResult> {
+  // Fast path: if rawBase64 is present on both, identical, and no cropping needed,
+  // skip all decoding entirely.
+  if (
+    firstScreenshot.rawBase64 &&
+    secondScreenshot.rawBase64 &&
+    firstScreenshot.rawBase64 === secondScreenshot.rawBase64 &&
+    !firstScreenshot.offset &&
+    !secondScreenshot.offset
+  ) {
+    const buffer = Buffer.from(firstScreenshot.rawBase64, 'base64');
+    return { firstImage: buffer, secondImage: buffer, diffImage: null, isEqual: true, diffPixels: 0 };
+  }
+
   const pixelRatio = firstScreenshot.pixelRatio || 1;
   const size = normalizeSize(firstScreenshot, secondScreenshot);
-  const firstImageCropRect: ElementRect = {
-    height: size.height,
-    width: size.width,
-    bottom: firstScreenshot.offset.top + size.height,
-    right: firstScreenshot.offset.left + size.width,
-    top: firstScreenshot.offset.top,
-    left: firstScreenshot.offset.left,
-  };
-  const secondImageCropRect: ElementRect = {
-    height: size.height,
-    width: size.width,
-    bottom: secondScreenshot.offset.top + size.height,
-    right: secondScreenshot.offset.left + size.width,
-    top: secondScreenshot.offset.top,
-    left: secondScreenshot.offset.left,
-  };
-  const firstImage = cropImage(firstScreenshot.image, firstImageCropRect, pixelRatio);
-  const secondImage = cropImage(secondScreenshot.image, secondImageCropRect, pixelRatio);
-  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, scaleSize(size, pixelRatio));
+  const scaledSize = scaleSize(size, pixelRatio);
+
+  const firstNeedsCrop = !!firstScreenshot.offset;
+  const secondNeedsCrop = !!secondScreenshot.offset;
+
+  // Decode images on demand: use pre-decoded image if available, otherwise parse from rawBase64
+  const firstDecoded = firstScreenshot.image ?? (await parsePng(firstScreenshot.rawBase64));
+  const secondDecoded = secondScreenshot.image ?? (await parsePng(secondScreenshot.rawBase64));
+
+  const firstImage = firstNeedsCrop
+    ? cropImage(firstDecoded, buildCropRect(firstScreenshot as ScreenshotWithOffset, size), pixelRatio)
+    : firstDecoded;
+  const secondImage = secondNeedsCrop
+    ? cropImage(secondDecoded, buildCropRect(secondScreenshot as ScreenshotWithOffset, size), pixelRatio)
+    : secondDecoded;
+
+  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, scaledSize);
+
+  // Skip packPng when no cropping was needed and rawBase64 is available
   const [firstPacked, secondPacked, diffPacked] = await Promise.all([
-    packPng(firstImage),
-    packPng(secondImage),
+    !firstNeedsCrop && firstScreenshot.rawBase64
+      ? Buffer.from(firstScreenshot.rawBase64, 'base64')
+      : packPng(firstImage),
+    !secondNeedsCrop && secondScreenshot.rawBase64
+      ? Buffer.from(secondScreenshot.rawBase64, 'base64')
+      : packPng(secondImage),
     diffImage && packPng(diffImage),
   ]);
+
   return {
     firstImage: firstPacked,
     secondImage: secondPacked,
@@ -79,3 +101,14 @@ export async function cropAndCompare(
     diffPixels,
   };
 }
+
+function buildCropRect(screenshot: ScreenshotWithOffset, size: ElementSize): ElementRect {
+  return {
+    height: size.height,
+    width: size.width,
+    bottom: screenshot.offset.top + size.height,
+    right: screenshot.offset.left + size.width,
+    top: screenshot.offset.top,
+    left: screenshot.offset.left,
+  };
+}
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 31c4691..fb227cb 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -9,10 +9,11 @@ import {
 } from '../browser-scripts';
 import { parsePng } from '../image-utils';
 import BasePageObject from './base';
-import { ElementOffset, ScreenshotCapturingOptions, ScreenshotWithOffset } from './types';
+import { ElementOffset, Screenshot, ScreenshotCapturingOptions } from './types';
 import fullPageScreenshot from './full-page-screenshot';
 
-export interface PermutationScreenshot extends ScreenshotWithOffset {
+export interface PermutationScreenshot extends Screenshot {
+  /** Identifier from the data-permutation attribute */
   id: string;
 }
 
@@ -46,54 +47,95 @@ export default class ScreenshotPageObject extends BasePageObject {
     return screenshot;
   }
 
-  async captureBySelector(selector: string, options: ScreenshotCapturingOptions = {}): Promise<ScreenshotWithOffset> {
+  async captureBySelector(selector: string, options: ScreenshotCapturingOptions = {}): Promise<Screenshot> {
     await this.waitForVisible(selector);
-    const { pixelRatio, top, left } = await this.getViewportSize();
+    const { pixelRatio } = await this.getViewportSize();
     const box = await this.getBoundingBox(selector);
-    const screenshot = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
-    const image = await parsePng(screenshot);
-
-    const offset: ElementOffset = { top: box.top, left: box.left };
-    if (!options.viewportOnly) {
-      // Correct potential scrolling offsets when using a full page screenshot
-      offset.top += top;
-      offset.left += left;
-    }
 
-    return { image, offset, pixelRatio, height: box.height, width: box.width };
+    try {
+      // Fast path: element screenshot is already cropped, no parsePng needed
+      const element = this.browser.$(selector);
+      const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
+      return { rawBase64, pixelRatio, height: box.height, width: box.width };
+    } catch {
+      // Fallback: full-page or viewport screenshot with offset
+      const { top, left } = await this.getViewportSize();
+      const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
+      const image = await parsePng(rawBase64);
+
+      const offset: ElementOffset = { top: box.top, left: box.left };
+      if (!options.viewportOnly) {
+        offset.top += top;
+        offset.left += left;
+      }
+
+      return { image, offset, pixelRatio, height: box.height, width: box.width, rawBase64 };
+    }
   }
 
-  async captureViewport(): Promise<ScreenshotWithOffset> {
+  async captureViewport(): Promise<Screenshot> {
     const { height, width } = await this.getViewportSize();
-
-    const offset: ElementOffset = {
-      top: 0,
-      left: 0,
-    };
-
-    const screenshot = await this.browser.takeScreenshot();
-    const image = await parsePng(screenshot);
-    return { image, offset, height, width };
+    const rawBase64 = await this.browser.takeScreenshot();
+    return { height, width, rawBase64 };
   }
 
+  /**
+   * Captures all permutation elements. Uses takeElementScreenshot when available
+   * to return pre-cropped PNGs without decoding. Falls back to a single full-page
+   * screenshot with bounding box metadata if takeElementScreenshot is unavailable.
+   *
+   * Consumers can compare rawBase64 directly for fast byte-equality checks and
+   * only call parsePng(rawBase64) when a diff is needed.
+   */
   async capturePermutations(): Promise<PermutationScreenshot[]> {
     await this.windowScrollTo({ top: 0, left: 0 });
 
-    // Adapt viewport height to fit all elements before taking a screenshot
-    const originalWindowSize = await this.fitWindowHeightToContent();
+    const elements = await this.browser.$$('[data-permutation]');
 
-    const screenshot = await this.fullPageScreenshot();
-    const image = await parsePng(screenshot);
-    const permutations = await this.browser.execute(getPermutationSizes);
+    if ((await elements.length) === 0) {
+      throw new Error('No permutations found on current page.');
+    }
 
-    // Restore window size after taking the screenshot
-    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+    // Adapt viewport height to fit all elements before taking screenshots
+    const originalWindowSize = await this.fitWindowHeightToContent();
 
-    if (permutations.length === 0) {
-      throw new Error('No permutations found on current page.');
+    let results: PermutationScreenshot[];
+    try {
+      // Fast path: capture each element individually via takeElementScreenshot.
+      // Each screenshot is already cropped to the element — no parsePng needed.
+      const pixelRatio = await this.browser.execute(function () {
+        return window.devicePixelRatio || 1;
+      });
+      results = [];
+      for (const element of elements) {
+        const id = (await element.getAttribute('data-permutation')) || '';
+        const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
+        const size = await element.getSize();
+        results.push({
+          id,
+          rawBase64,
+          width: size.width * pixelRatio,
+          height: size.height * pixelRatio,
+        });
+      }
+    } catch {
+      // Fallback: single full-page screenshot with bounding box metadata
+      const permutations = await this.browser.execute(getPermutationSizes);
+      const rawBase64 = await this.browser.takeScreenshot();
+      const image = await parsePng(rawBase64);
+
+      results = permutations.map((permutation: PermutationInfo) => ({
+        id: permutation.id,
+        rawBase64,
+        image,
+        offset: permutation.offset,
+        width: permutation.width,
+        height: permutation.height,
+      }));
     }
 
-    return permutations.map((permutation: PermutationInfo) => ({ ...permutation, image }));
+    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+    return results;
   }
 
   private async fitWindowHeightToContent(): Promise<{ width: number; height: number }> {
diff --git a/src/page-objects/types.ts b/src/page-objects/types.ts
index 658df92..95b4b84 100644
--- a/src/page-objects/types.ts
+++ b/src/page-objects/types.ts
@@ -28,8 +28,18 @@ export interface ScreenshotCapturingOptions {
   viewportOnly?: boolean;
 }
 
-export interface ScreenshotWithOffset extends ElementSize {
-  image: PNG;
-  offset: ElementOffset;
+export interface Screenshot extends ElementSize {
+  image?: PNG;
   pixelRatio?: number;
+  /**
+   * The raw base64-encoded PNG from WebDriver, retained for fast byte-equality
+   * comparison in cropAndCompare. When two screenshots have the same rawBase64
+   * and no cropping is needed, expensive decoding is skipped entirely.
+   */
+  rawBase64: string;
+  offset?: ElementOffset;
+}
+
+export interface ScreenshotWithOffset extends Screenshot {
+  offset: ElementOffset;
 }
diff --git a/test/page-object.test.ts b/test/page-object.test.ts
index 204a4e2..ff099bd 100644
--- a/test/page-object.test.ts
+++ b/test/page-object.test.ts
@@ -373,16 +373,11 @@ describe('capturePermutations', () => {
       expect(permutations.map(p => p.id)).toEqual(['perm-1', 'perm-2', 'perm-3']);
 
       permutations.forEach(perm => {
-        expect(perm.image).toBeDefined();
+        expect(perm.rawBase64).toBeDefined();
+        expect(perm.rawBase64.length).toBeGreaterThan(0);
         expect(perm.width).toBeGreaterThan(0);
         expect(perm.height).toBeGreaterThan(0);
-        expect(perm.offset.top).toBeGreaterThanOrEqual(0);
-        expect(perm.offset.left).toBeGreaterThanOrEqual(0);
       });
-
-      // Verify permutations are in order (top to bottom)
-      expect(permutations[0].offset.top).toBeLessThan(permutations[1].offset.top);
-      expect(permutations[1].offset.top).toBeLessThan(permutations[2].offset.top);
     }, './test-permutations.html')
   );
 

From 130c87a2e06e8bd2bcc94c9204febe7f11aa11b3 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Fri, 19 Jun 2026 13:41:32 +0200
Subject: [PATCH 02/23] Export screenshot

---
 src/image-utils/compare.ts | 4 ----
 src/page-objects/index.ts  | 2 +-
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index b5be96e..b0d2aa3 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -33,10 +33,6 @@ function scaleSize(size: ElementSize, pixelRatio: number) {
   };
 }
 
-function isZeroOffset(screenshot: ScreenshotWithOffset): boolean {
-  return !screenshot.offset;
-}
-
 export interface CropAndCompareResult {
   firstImage: Buffer;
   secondImage: Buffer;
diff --git a/src/page-objects/index.ts b/src/page-objects/index.ts
index 43be3d8..3e937c4 100644
--- a/src/page-objects/index.ts
+++ b/src/page-objects/index.ts
@@ -3,4 +3,4 @@
 export { default as BasePageObject } from './base';
 export { default as ScreenshotPageObject, PermutationScreenshot } from './screenshot';
 export { default as EventsSpy } from './events-spy';
-export { ScreenshotWithOffset, ElementSize, ElementRect, ElementOffset } from './types';
+export { Screenshot, ScreenshotWithOffset, ElementSize, ElementRect, ElementOffset } from './types';

From 5ed8d1fe2bda4119f2fee812381818cb3b4451fc Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Sun, 21 Jun 2026 16:07:36 +0200
Subject: [PATCH 03/23] Make takeElementScreenshot opt-in

---
 src/page-objects/screenshot.ts | 66 +++++++++++++++++++---------------
 1 file changed, 37 insertions(+), 29 deletions(-)

diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index fb227cb..4836b97 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -87,44 +87,55 @@ export default class ScreenshotPageObject extends BasePageObject {
    * Consumers can compare rawBase64 directly for fast byte-equality checks and
    * only call parsePng(rawBase64) when a diff is needed.
    */
-  async capturePermutations(): Promise<PermutationScreenshot[]> {
-    await this.windowScrollTo({ top: 0, left: 0 });
+  async capturePermutations(options?: { individualScreenshots?: boolean }): Promise<PermutationScreenshot[]> {
+    // Adapt viewport height to fit all elements before taking screenshots
+    const originalWindowSize = await this.fitWindowHeightToContent();
+
+    const results = this.takePermutationScreenshots(options?.individualScreenshots);
+
+    // Restore original viewport height
+    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
 
-    const elements = await this.browser.$$('[data-permutation]');
+    return results;
+  }
 
+  private async takePermutationScreenshots(individualScreenshots = false) {
+    const elements = this.browser.$$('[data-permutation]');
     if ((await elements.length) === 0) {
       throw new Error('No permutations found on current page.');
     }
-
-    // Adapt viewport height to fit all elements before taking screenshots
-    const originalWindowSize = await this.fitWindowHeightToContent();
-
-    let results: PermutationScreenshot[];
-    try {
-      // Fast path: capture each element individually via takeElementScreenshot.
-      // Each screenshot is already cropped to the element — no parsePng needed.
-      const pixelRatio = await this.browser.execute(function () {
-        return window.devicePixelRatio || 1;
-      });
-      results = [];
-      for (const element of elements) {
-        const id = (await element.getAttribute('data-permutation')) || '';
-        const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
-        const size = await element.getSize();
-        results.push({
-          id,
-          rawBase64,
-          width: size.width * pixelRatio,
-          height: size.height * pixelRatio,
+    if (individualScreenshots) {
+      try {
+        // Try to capture each element individually via takeElementScreenshot.
+        // Each screenshot is already cropped to the element — no need to decode, crop and re-encode to PNG.
+        const results: PermutationScreenshot[] = [];
+        const pixelRatio = await this.browser.execute(function () {
+          return window.devicePixelRatio || 1;
         });
+        for (const element of elements) {
+          const id = (await element.getAttribute('data-permutation')) || '';
+          const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
+          const size = await element.getSize();
+          results.push({
+            id,
+            rawBase64,
+            width: size.width * pixelRatio,
+            height: size.height * pixelRatio,
+          });
+        }
+        return results;
+      } catch {
+        // If takeElementScreenshot fails (for example for browsers where this API is not available),
+        // fall back to taking one single screenshot and cropping it
+        return this.capturePermutations();
       }
-    } catch {
+    } else {
       // Fallback: single full-page screenshot with bounding box metadata
       const permutations = await this.browser.execute(getPermutationSizes);
       const rawBase64 = await this.browser.takeScreenshot();
       const image = await parsePng(rawBase64);
 
-      results = permutations.map((permutation: PermutationInfo) => ({
+      return permutations.map((permutation: PermutationInfo) => ({
         id: permutation.id,
         rawBase64,
         image,
@@ -133,9 +144,6 @@ export default class ScreenshotPageObject extends BasePageObject {
         height: permutation.height,
       }));
     }
-
-    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
-    return results;
   }
 
   private async fitWindowHeightToContent(): Promise<{ width: number; height: number }> {

From e0c534334ef938fd0a91b7a77b45bcd8cd3b4d4d Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Sun, 21 Jun 2026 21:10:17 +0200
Subject: [PATCH 04/23] Refinements

---
 src/image-utils/compare.ts     | 18 ++++++++++--------
 src/page-objects/screenshot.ts |  8 +++-----
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index b0d2aa3..9ef6035 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -3,7 +3,7 @@
 import { PNG } from 'pngjs';
 import pixelmatch from 'pixelmatch';
 import { packPng, cropImage, parsePng } from './utils';
-import { ElementRect, ElementSize, Screenshot, ScreenshotWithOffset } from '../page-objects/types';
+import { ElementRect, ElementSize, Screenshot } from '../page-objects/types';
 
 export function compareImages(firstImage: PNG, secondImage: PNG, { width, height }: ElementSize) {
   // This prevents an error thrown from pixelmatch when comparing 0-sized images.
@@ -70,10 +70,10 @@ export async function cropAndCompare(
   const secondDecoded = secondScreenshot.image ?? (await parsePng(secondScreenshot.rawBase64));
 
   const firstImage = firstNeedsCrop
-    ? cropImage(firstDecoded, buildCropRect(firstScreenshot as ScreenshotWithOffset, size), pixelRatio)
+    ? cropImage(firstDecoded, buildCropRect(firstScreenshot, size), pixelRatio)
     : firstDecoded;
   const secondImage = secondNeedsCrop
-    ? cropImage(secondDecoded, buildCropRect(secondScreenshot as ScreenshotWithOffset, size), pixelRatio)
+    ? cropImage(secondDecoded, buildCropRect(secondScreenshot, size), pixelRatio)
     : secondDecoded;
 
   const { diffImage, diffPixels } = compareImages(firstImage, secondImage, scaledSize);
@@ -98,13 +98,15 @@ export async function cropAndCompare(
   };
 }
 
-function buildCropRect(screenshot: ScreenshotWithOffset, size: ElementSize): ElementRect {
+function buildCropRect(screenshot: Screenshot, size: ElementSize): ElementRect {
+  const top = screenshot.offset?.top ?? 0;
+  const left = screenshot.offset?.left ?? 0;
   return {
     height: size.height,
     width: size.width,
-    bottom: screenshot.offset.top + size.height,
-    right: screenshot.offset.left + size.width,
-    top: screenshot.offset.top,
-    left: screenshot.offset.left,
+    bottom: top + size.height,
+    right: left + size.width,
+    top,
+    left,
   };
 }
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 4836b97..cf61a5d 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -99,7 +99,7 @@ export default class ScreenshotPageObject extends BasePageObject {
     return results;
   }
 
-  private async takePermutationScreenshots(individualScreenshots = false) {
+  private async takePermutationScreenshots(individualScreenshots = false): Promise<PermutationScreenshot[]> {
     const elements = this.browser.$$('[data-permutation]');
     if ((await elements.length) === 0) {
       throw new Error('No permutations found on current page.');
@@ -127,18 +127,16 @@ export default class ScreenshotPageObject extends BasePageObject {
       } catch {
         // If takeElementScreenshot fails (for example for browsers where this API is not available),
         // fall back to taking one single screenshot and cropping it
-        return this.capturePermutations();
+        return this.takePermutationScreenshots();
       }
     } else {
-      // Fallback: single full-page screenshot with bounding box metadata
+      // Single full-page screenshot with bounding box metadata
       const permutations = await this.browser.execute(getPermutationSizes);
       const rawBase64 = await this.browser.takeScreenshot();
-      const image = await parsePng(rawBase64);
 
       return permutations.map((permutation: PermutationInfo) => ({
         id: permutation.id,
         rawBase64,
-        image,
         offset: permutation.offset,
         width: permutation.width,
         height: permutation.height,

From 268b39fcb0a87a7d3784c4405e7641182ec4c9be Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Sun, 21 Jun 2026 21:33:39 +0200
Subject: [PATCH 05/23] Use fullPageScreenshot for fallback permutation capture
 method

---
 src/page-objects/screenshot.ts | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index cf61a5d..66fd49d 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -93,7 +93,7 @@ export default class ScreenshotPageObject extends BasePageObject {
 
     const results = this.takePermutationScreenshots(options?.individualScreenshots);
 
-    // Restore original viewport height
+    // Restore window size after taking the screenshot
     await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
 
     return results;
@@ -130,13 +130,14 @@ export default class ScreenshotPageObject extends BasePageObject {
         return this.takePermutationScreenshots();
       }
     } else {
-      // Single full-page screenshot with bounding box metadata
+      // Single full-page screenshot with bounding box metadata for cropping
+
+      const screenshot = await this.fullPageScreenshot();
       const permutations = await this.browser.execute(getPermutationSizes);
-      const rawBase64 = await this.browser.takeScreenshot();
 
       return permutations.map((permutation: PermutationInfo) => ({
         id: permutation.id,
-        rawBase64,
+        rawBase64: screenshot,
         offset: permutation.offset,
         width: permutation.width,
         height: permutation.height,

From 8ad4cf9f58ee8252b43b37db2e15bd328dca73b9 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Mon, 22 Jun 2026 06:42:01 +0200
Subject: [PATCH 06/23] Scroll to top before capturing permutations

---
 src/page-objects/screenshot.ts | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 66fd49d..26da158 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -53,12 +53,11 @@ export default class ScreenshotPageObject extends BasePageObject {
     const box = await this.getBoundingBox(selector);
 
     try {
-      // Fast path: element screenshot is already cropped, no parsePng needed
       const element = this.browser.$(selector);
       const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
       return { rawBase64, pixelRatio, height: box.height, width: box.width };
     } catch {
-      // Fallback: full-page or viewport screenshot with offset
+      console.warn('Could not use takeElementScreenshot. Falling back to full-page screenshot and cropping');
       const { top, left } = await this.getViewportSize();
       const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
       const image = await parsePng(rawBase64);
@@ -88,6 +87,8 @@ export default class ScreenshotPageObject extends BasePageObject {
    * only call parsePng(rawBase64) when a diff is needed.
    */
   async capturePermutations(options?: { individualScreenshots?: boolean }): Promise<PermutationScreenshot[]> {
+    await this.windowScrollTo({ top: 0, left: 0 });
+
     // Adapt viewport height to fit all elements before taking screenshots
     const originalWindowSize = await this.fitWindowHeightToContent();
 
@@ -125,8 +126,7 @@ export default class ScreenshotPageObject extends BasePageObject {
         }
         return results;
       } catch {
-        // If takeElementScreenshot fails (for example for browsers where this API is not available),
-        // fall back to taking one single screenshot and cropping it
+        console.warn('Could not use takeElementScreenshot. Falling back to full-page screenshot and cropping');
         return this.takePermutationScreenshots();
       }
     } else {

From cf21b5e2c67a211431b10851f3a2bbf7fc500e2d Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Mon, 22 Jun 2026 21:24:55 +0200
Subject: [PATCH 07/23] Refactor

---
 src/image-utils/compare.ts | 60 ++++++++++++++++++++++++++------------
 1 file changed, 42 insertions(+), 18 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index 9ef6035..32c2b4c 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -3,7 +3,12 @@
 import { PNG } from 'pngjs';
 import pixelmatch from 'pixelmatch';
 import { packPng, cropImage, parsePng } from './utils';
-import { ElementRect, ElementSize, Screenshot } from '../page-objects/types';
+import { ElementRect, ElementSize, Screenshot, ScreenshotWithOffset } from '../page-objects/types';
+
+interface Size {
+  width: number;
+  height: number;
+}
 
 export function compareImages(firstImage: PNG, secondImage: PNG, { width, height }: ElementSize) {
   // This prevents an error thrown from pixelmatch when comparing 0-sized images.
@@ -19,14 +24,14 @@ export function compareImages(firstImage: PNG, secondImage: PNG, { width, height
   return { diffPixels, diffImage };
 }
 
-function normalizeSize(firstScreenshot: Screenshot, secondScreenshot: Screenshot) {
+function normalizeSize(firstScreenshot: Screenshot, secondScreenshot: Screenshot): Size {
   return {
     height: Math.round(Math.max(firstScreenshot.height, secondScreenshot.height)),
     width: Math.round(Math.max(firstScreenshot.width, secondScreenshot.width)),
   };
 }
 
-function scaleSize(size: ElementSize, pixelRatio: number) {
+function scaleSize(size: ElementSize, pixelRatio: number): Size {
   return {
     width: Math.ceil(size.width * pixelRatio),
     height: Math.ceil(size.height * pixelRatio),
@@ -41,6 +46,22 @@ export interface CropAndCompareResult {
   diffPixels: number;
 }
 
+async function getScreenshotImage(screenshot: Screenshot) {
+  if (!screenshot.image) {
+    screenshot.image = await parsePng(screenshot.rawBase64);
+  }
+  return screenshot.image;
+}
+
+async function cropIfNeeded(screenshot: Screenshot, size: Size) {
+  const image = await getScreenshotImage(screenshot);
+  if (screenshot.offset) {
+    return cropImage(image, buildCropRect(screenshot as ScreenshotWithOffset, size), screenshot.pixelRatio);
+  } else {
+    return image;
+  }
+}
+
 export async function cropAndCompare(
   firstScreenshot: Screenshot,
   secondScreenshot: Screenshot
@@ -58,32 +79,35 @@ export async function cropAndCompare(
     return { firstImage: buffer, secondImage: buffer, diffImage: null, isEqual: true, diffPixels: 0 };
   }
 
+  console.log('first', { ...firstScreenshot, image: undefined, rawBase64: undefined });
+  console.log('second', { ...secondScreenshot, image: undefined, rawBase64: undefined });
+
   const pixelRatio = firstScreenshot.pixelRatio || 1;
+
   const size = normalizeSize(firstScreenshot, secondScreenshot);
   const scaledSize = scaleSize(size, pixelRatio);
 
-  const firstNeedsCrop = !!firstScreenshot.offset;
-  const secondNeedsCrop = !!secondScreenshot.offset;
+  console.log({ scaledSize });
+
+  const firstImage = await cropIfNeeded(firstScreenshot, scaledSize);
+  const secondImage = await cropIfNeeded(secondScreenshot, scaledSize);
 
-  // Decode images on demand: use pre-decoded image if available, otherwise parse from rawBase64
-  const firstDecoded = firstScreenshot.image ?? (await parsePng(firstScreenshot.rawBase64));
-  const secondDecoded = secondScreenshot.image ?? (await parsePng(secondScreenshot.rawBase64));
+  // Make sure the size of both images to compare is the same
+  const compareSize = {
+    width: Math.max(firstImage.width, secondImage.width),
+    height: Math.max(firstImage.height, secondImage.height),
+  };
 
-  const firstImage = firstNeedsCrop
-    ? cropImage(firstDecoded, buildCropRect(firstScreenshot, size), pixelRatio)
-    : firstDecoded;
-  const secondImage = secondNeedsCrop
-    ? cropImage(secondDecoded, buildCropRect(secondScreenshot, size), pixelRatio)
-    : secondDecoded;
+  console.log({ compareSize });
 
-  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, scaledSize);
+  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, compareSize);
 
   // Skip packPng when no cropping was needed and rawBase64 is available
   const [firstPacked, secondPacked, diffPacked] = await Promise.all([
-    !firstNeedsCrop && firstScreenshot.rawBase64
+    !firstScreenshot.offset && firstScreenshot.rawBase64
       ? Buffer.from(firstScreenshot.rawBase64, 'base64')
       : packPng(firstImage),
-    !secondNeedsCrop && secondScreenshot.rawBase64
+    !secondScreenshot.offset && secondScreenshot.rawBase64
       ? Buffer.from(secondScreenshot.rawBase64, 'base64')
       : packPng(secondImage),
     diffImage && packPng(diffImage),
@@ -98,7 +122,7 @@ export async function cropAndCompare(
   };
 }
 
-function buildCropRect(screenshot: Screenshot, size: ElementSize): ElementRect {
+function buildCropRect(screenshot: ScreenshotWithOffset, size: ElementSize): ElementRect {
   const top = screenshot.offset?.top ?? 0;
   const left = screenshot.offset?.left ?? 0;
   return {

From 9f991ecf1932ff96478a6ccec10b19383fb2934f Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Tue, 23 Jun 2026 07:40:23 +0200
Subject: [PATCH 08/23] Always let capturePermutations return parsed PNG when
 not taking indivudual screenshots

---
 src/page-objects/screenshot.ts | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 26da158..c8f2d3a 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -133,15 +133,10 @@ export default class ScreenshotPageObject extends BasePageObject {
       // Single full-page screenshot with bounding box metadata for cropping
 
       const screenshot = await this.fullPageScreenshot();
+      const image = await parsePng(screenshot);
       const permutations = await this.browser.execute(getPermutationSizes);
 
-      return permutations.map((permutation: PermutationInfo) => ({
-        id: permutation.id,
-        rawBase64: screenshot,
-        offset: permutation.offset,
-        width: permutation.width,
-        height: permutation.height,
-      }));
+      return permutations.map((permutation: PermutationInfo) => ({ ...permutation, image, rawBase64: screenshot }));
     }
   }
 

From 61bdd9c01930ea6067e56e3abfd7f4770f6cbaa2 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Tue, 23 Jun 2026 07:52:16 +0200
Subject: [PATCH 09/23] Make page tall as needed in captureBySelector

---
 src/image-utils/compare.ts     | 26 +++-----------------------
 src/page-objects/screenshot.ts |  6 ++++++
 2 files changed, 9 insertions(+), 23 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index 32c2b4c..4bf1c5f 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -66,41 +66,21 @@ export async function cropAndCompare(
   firstScreenshot: Screenshot,
   secondScreenshot: Screenshot
 ): Promise<CropAndCompareResult> {
-  // Fast path: if rawBase64 is present on both, identical, and no cropping needed,
-  // skip all decoding entirely.
-  if (
-    firstScreenshot.rawBase64 &&
-    secondScreenshot.rawBase64 &&
-    firstScreenshot.rawBase64 === secondScreenshot.rawBase64 &&
-    !firstScreenshot.offset &&
-    !secondScreenshot.offset
-  ) {
+  // Fast path: if rawBase64 is present on both, identical, and no cropping needed, skip all decoding entirely.
+  if (!firstScreenshot.offset && !secondScreenshot.offset && firstScreenshot.rawBase64 === secondScreenshot.rawBase64) {
     const buffer = Buffer.from(firstScreenshot.rawBase64, 'base64');
     return { firstImage: buffer, secondImage: buffer, diffImage: null, isEqual: true, diffPixels: 0 };
   }
 
-  console.log('first', { ...firstScreenshot, image: undefined, rawBase64: undefined });
-  console.log('second', { ...secondScreenshot, image: undefined, rawBase64: undefined });
-
   const pixelRatio = firstScreenshot.pixelRatio || 1;
 
   const size = normalizeSize(firstScreenshot, secondScreenshot);
   const scaledSize = scaleSize(size, pixelRatio);
 
-  console.log({ scaledSize });
-
   const firstImage = await cropIfNeeded(firstScreenshot, scaledSize);
   const secondImage = await cropIfNeeded(secondScreenshot, scaledSize);
 
-  // Make sure the size of both images to compare is the same
-  const compareSize = {
-    width: Math.max(firstImage.width, secondImage.width),
-    height: Math.max(firstImage.height, secondImage.height),
-  };
-
-  console.log({ compareSize });
-
-  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, compareSize);
+  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, size);
 
   // Skip packPng when no cropping was needed and rawBase64 is available
   const [firstPacked, secondPacked, diffPacked] = await Promise.all([
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index c8f2d3a..a84dedc 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -53,11 +53,17 @@ export default class ScreenshotPageObject extends BasePageObject {
     const box = await this.getBoundingBox(selector);
 
     try {
+      const originalWindowSize = await this.fitWindowHeightToContent();
+
       const element = this.browser.$(selector);
       const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
+
+      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+
       return { rawBase64, pixelRatio, height: box.height, width: box.width };
     } catch {
       console.warn('Could not use takeElementScreenshot. Falling back to full-page screenshot and cropping');
+
       const { top, left } = await this.getViewportSize();
       const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
       const image = await parsePng(rawBase64);

From 83ad3a6b8613b5fe1486006682a103692b887dca Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Tue, 23 Jun 2026 10:53:34 +0200
Subject: [PATCH 10/23] Make takeElementScreenshot opt-in also in
 captureBySelector

---
 src/page-objects/screenshot.ts | 37 +++++++++++++++-------------------
 src/page-objects/types.ts      |  1 +
 2 files changed, 17 insertions(+), 21 deletions(-)

diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index a84dedc..9025ea6 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -52,18 +52,21 @@ export default class ScreenshotPageObject extends BasePageObject {
     const { pixelRatio } = await this.getViewportSize();
     const box = await this.getBoundingBox(selector);
 
-    try {
-      const originalWindowSize = await this.fitWindowHeightToContent();
-
-      const element = this.browser.$(selector);
-      const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
+    if (options?.singleElements) {
+      try {
+        const originalWindowSize = await this.fitWindowHeightToContent();
 
-      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+        const element = this.browser.$(selector);
+        const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
 
-      return { rawBase64, pixelRatio, height: box.height, width: box.width };
-    } catch {
-      console.warn('Could not use takeElementScreenshot. Falling back to full-page screenshot and cropping');
+        await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
 
+        return { rawBase64, pixelRatio, height: box.height, width: box.width };
+      } catch {
+        console.warn('Could not use takeElementScreenshot. Falling back to full-page screenshot and cropping');
+        return this.captureBySelector(selector, { ...options, singleElements: false });
+      }
+    } else {
       const { top, left } = await this.getViewportSize();
       const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
       const image = await parsePng(rawBase64);
@@ -84,21 +87,13 @@ export default class ScreenshotPageObject extends BasePageObject {
     return { height, width, rawBase64 };
   }
 
-  /**
-   * Captures all permutation elements. Uses takeElementScreenshot when available
-   * to return pre-cropped PNGs without decoding. Falls back to a single full-page
-   * screenshot with bounding box metadata if takeElementScreenshot is unavailable.
-   *
-   * Consumers can compare rawBase64 directly for fast byte-equality checks and
-   * only call parsePng(rawBase64) when a diff is needed.
-   */
-  async capturePermutations(options?: { individualScreenshots?: boolean }): Promise<PermutationScreenshot[]> {
+  async capturePermutations(options?: { singleElements?: boolean }): Promise<PermutationScreenshot[]> {
     await this.windowScrollTo({ top: 0, left: 0 });
 
     // Adapt viewport height to fit all elements before taking screenshots
     const originalWindowSize = await this.fitWindowHeightToContent();
 
-    const results = this.takePermutationScreenshots(options?.individualScreenshots);
+    const results = this.takePermutationScreenshots(options?.singleElements);
 
     // Restore window size after taking the screenshot
     await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
@@ -106,12 +101,12 @@ export default class ScreenshotPageObject extends BasePageObject {
     return results;
   }
 
-  private async takePermutationScreenshots(individualScreenshots = false): Promise<PermutationScreenshot[]> {
+  private async takePermutationScreenshots(singleElements = false): Promise<PermutationScreenshot[]> {
     const elements = this.browser.$$('[data-permutation]');
     if ((await elements.length) === 0) {
       throw new Error('No permutations found on current page.');
     }
-    if (individualScreenshots) {
+    if (singleElements) {
       try {
         // Try to capture each element individually via takeElementScreenshot.
         // Each screenshot is already cropped to the element — no need to decode, crop and re-encode to PNG.
diff --git a/src/page-objects/types.ts b/src/page-objects/types.ts
index 95b4b84..6840f5a 100644
--- a/src/page-objects/types.ts
+++ b/src/page-objects/types.ts
@@ -26,6 +26,7 @@ export interface ViewportSize extends ElementOffset, ElementSize {
 
 export interface ScreenshotCapturingOptions {
   viewportOnly?: boolean;
+  singleElements?: boolean;
 }
 
 export interface Screenshot extends ElementSize {

From e2d117545eef1441f11fdcc9fab8aaca56782873 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Wed, 24 Jun 2026 00:30:18 +0200
Subject: [PATCH 11/23] Make implementation backwards-compatible

---
 src/image-utils/compare.ts     |  18 ++--
 src/page-objects/index.ts      |   9 +-
 src/page-objects/screenshot.ts | 157 +++++++++++++++++++--------------
 src/page-objects/types.ts      |  24 +++--
 4 files changed, 123 insertions(+), 85 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index 4bf1c5f..15c9c00 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -31,6 +31,10 @@ function normalizeSize(firstScreenshot: Screenshot, secondScreenshot: Screenshot
   };
 }
 
+function isScreenshotWithOffset(s: Screenshot): s is ScreenshotWithOffset {
+  return 'offset' in s && s.offset !== undefined;
+}
+
 function scaleSize(size: ElementSize, pixelRatio: number): Size {
   return {
     width: Math.ceil(size.width * pixelRatio),
@@ -55,7 +59,7 @@ async function getScreenshotImage(screenshot: Screenshot) {
 
 async function cropIfNeeded(screenshot: Screenshot, size: Size) {
   const image = await getScreenshotImage(screenshot);
-  if (screenshot.offset) {
+  if (isScreenshotWithOffset(screenshot)) {
     return cropImage(image, buildCropRect(screenshot as ScreenshotWithOffset, size), screenshot.pixelRatio);
   } else {
     return image;
@@ -67,7 +71,11 @@ export async function cropAndCompare(
   secondScreenshot: Screenshot
 ): Promise<CropAndCompareResult> {
   // Fast path: if rawBase64 is present on both, identical, and no cropping needed, skip all decoding entirely.
-  if (!firstScreenshot.offset && !secondScreenshot.offset && firstScreenshot.rawBase64 === secondScreenshot.rawBase64) {
+  if (
+    !isScreenshotWithOffset(firstScreenshot) &&
+    !isScreenshotWithOffset(secondScreenshot) &&
+    firstScreenshot.rawBase64 === secondScreenshot.rawBase64
+  ) {
     const buffer = Buffer.from(firstScreenshot.rawBase64, 'base64');
     return { firstImage: buffer, secondImage: buffer, diffImage: null, isEqual: true, diffPixels: 0 };
   }
@@ -84,10 +92,8 @@ export async function cropAndCompare(
 
   // Skip packPng when no cropping was needed and rawBase64 is available
   const [firstPacked, secondPacked, diffPacked] = await Promise.all([
-    !firstScreenshot.offset && firstScreenshot.rawBase64
-      ? Buffer.from(firstScreenshot.rawBase64, 'base64')
-      : packPng(firstImage),
-    !secondScreenshot.offset && secondScreenshot.rawBase64
+    !isScreenshotWithOffset(firstScreenshot) ? Buffer.from(firstScreenshot.rawBase64, 'base64') : packPng(firstImage),
+    !isScreenshotWithOffset(secondScreenshot)
       ? Buffer.from(secondScreenshot.rawBase64, 'base64')
       : packPng(secondImage),
     diffImage && packPng(diffImage),
diff --git a/src/page-objects/index.ts b/src/page-objects/index.ts
index 3e937c4..da26948 100644
--- a/src/page-objects/index.ts
+++ b/src/page-objects/index.ts
@@ -1,6 +1,11 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 export { default as BasePageObject } from './base';
-export { default as ScreenshotPageObject, PermutationScreenshot } from './screenshot';
+export {
+  default as ScreenshotPageObject,
+  PermutationScreenshot,
+  RawPermutationScreenshot,
+  DecodedPermutationScreenshot,
+} from './screenshot';
 export { default as EventsSpy } from './events-spy';
-export { Screenshot, ScreenshotWithOffset, ElementSize, ElementRect, ElementOffset } from './types';
+export { Screenshot, RawScreenshot, ScreenshotWithOffset, ElementSize, ElementRect, ElementOffset } from './types';
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 9025ea6..880903a 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -9,14 +9,21 @@ import {
 } from '../browser-scripts';
 import { parsePng } from '../image-utils';
 import BasePageObject from './base';
-import { ElementOffset, Screenshot, ScreenshotCapturingOptions } from './types';
+import { ElementOffset, RawScreenshot, ScreenshotWithOffset, ScreenshotCapturingOptions } from './types';
 import fullPageScreenshot from './full-page-screenshot';
 
-export interface PermutationScreenshot extends Screenshot {
-  /** Identifier from the data-permutation attribute */
+/** Permutation screenshot from takeElementScreenshot — no image, no offset. */
+export interface RawPermutationScreenshot extends RawScreenshot {
   id: string;
 }
 
+/** Permutation screenshot from full-page fallback — has image and offset for cropping. */
+export interface DecodedPermutationScreenshot extends ScreenshotWithOffset {
+  id: string;
+}
+
+export type PermutationScreenshot = RawPermutationScreenshot | DecodedPermutationScreenshot;
+
 export default class ScreenshotPageObject extends BasePageObject {
   constructor(browser: WebdriverIO.Browser, public readonly forceScrollAndMerge: boolean = false) {
     super(browser);
@@ -47,97 +54,111 @@ export default class ScreenshotPageObject extends BasePageObject {
     return screenshot;
   }
 
-  async captureBySelector(selector: string, options: ScreenshotCapturingOptions = {}): Promise<Screenshot> {
+  // Overloads for captureBySelector
+  async captureBySelector(
+    selector: string,
+    options: ScreenshotCapturingOptions & { singleElements: true }
+  ): Promise<RawScreenshot>;
+  async captureBySelector(
+    selector: string,
+    options?: ScreenshotCapturingOptions & { singleElements?: false }
+  ): Promise<ScreenshotWithOffset>;
+  async captureBySelector(
+    selector: string,
+    options: ScreenshotCapturingOptions = {}
+  ): Promise<RawScreenshot | ScreenshotWithOffset> {
     await this.waitForVisible(selector);
     const { pixelRatio } = await this.getViewportSize();
     const box = await this.getBoundingBox(selector);
 
-    if (options?.singleElements) {
-      try {
-        const originalWindowSize = await this.fitWindowHeightToContent();
-
-        const element = this.browser.$(selector);
-        const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
-
-        await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
-
-        return { rawBase64, pixelRatio, height: box.height, width: box.width };
-      } catch {
-        console.warn('Could not use takeElementScreenshot. Falling back to full-page screenshot and cropping');
-        return this.captureBySelector(selector, { ...options, singleElements: false });
-      }
-    } else {
-      const { top, left } = await this.getViewportSize();
-      const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
-      const image = await parsePng(rawBase64);
+    if (options.singleElements) {
+      const originalWindowSize = await this.fitWindowHeightToContent();
+      const element = this.browser.$(selector);
+      const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
+      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+      return { rawBase64, pixelRatio, height: box.height, width: box.width };
+    }
 
-      const offset: ElementOffset = { top: box.top, left: box.left };
-      if (!options.viewportOnly) {
-        offset.top += top;
-        offset.left += left;
-      }
+    // Default: full-page or viewport screenshot with decoded image
+    const { top, left } = await this.getViewportSize();
+    const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
+    const image = await parsePng(rawBase64);
 
-      return { image, offset, pixelRatio, height: box.height, width: box.width, rawBase64 };
+    const offset: ElementOffset = { top: box.top, left: box.left };
+    if (!options.viewportOnly) {
+      offset.top += top;
+      offset.left += left;
     }
+
+    return { image, offset, pixelRatio, height: box.height, width: box.width, rawBase64 };
   }
 
-  async captureViewport(): Promise<Screenshot> {
+  async captureViewport(): Promise<ScreenshotWithOffset> {
     const { height, width } = await this.getViewportSize();
     const rawBase64 = await this.browser.takeScreenshot();
-    return { height, width, rawBase64 };
+    const image = await parsePng(rawBase64);
+    return { image, offset: { top: 0, left: 0 }, height, width, rawBase64 };
   }
 
+  // Overloads for capturePermutations
+  async capturePermutations(options: { singleElements: true }): Promise<RawPermutationScreenshot[]>;
+  async capturePermutations(options?: { singleElements?: false }): Promise<DecodedPermutationScreenshot[]>;
   async capturePermutations(options?: { singleElements?: boolean }): Promise<PermutationScreenshot[]> {
     await this.windowScrollTo({ top: 0, left: 0 });
 
     // Adapt viewport height to fit all elements before taking screenshots
     const originalWindowSize = await this.fitWindowHeightToContent();
 
-    const results = this.takePermutationScreenshots(options?.singleElements);
-
-    // Restore window size after taking the screenshot
-    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
-
-    return results;
+    try {
+      const results = await this.takePermutationScreenshots(options);
+      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+      return results;
+    } catch (error) {
+      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+      throw error;
+    }
   }
 
-  private async takePermutationScreenshots(singleElements = false): Promise<PermutationScreenshot[]> {
-    const elements = this.browser.$$('[data-permutation]');
-    if ((await elements.length) === 0) {
-      throw new Error('No permutations found on current page.');
-    }
-    if (singleElements) {
-      try {
-        // Try to capture each element individually via takeElementScreenshot.
-        // Each screenshot is already cropped to the element — no need to decode, crop and re-encode to PNG.
-        const results: PermutationScreenshot[] = [];
-        const pixelRatio = await this.browser.execute(function () {
-          return window.devicePixelRatio || 1;
+  private async takePermutationScreenshots(options?: { singleElements?: boolean }): Promise<PermutationScreenshot[]> {
+    if (options?.singleElements) {
+      const elements = this.browser.$$('[data-permutation]');
+      if ((await elements.length) === 0) {
+        throw new Error('No permutations found on current page.');
+      }
+
+      const pixelRatio = await this.browser.execute(function () {
+        return window.devicePixelRatio || 1;
+      });
+      const results: RawPermutationScreenshot[] = [];
+      for (const element of elements) {
+        const id = (await element.getAttribute('data-permutation')) || '';
+        const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
+        const size = await element.getSize();
+        results.push({
+          id,
+          rawBase64,
+          width: size.width * pixelRatio,
+          height: size.height * pixelRatio,
         });
-        for (const element of elements) {
-          const id = (await element.getAttribute('data-permutation')) || '';
-          const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
-          const size = await element.getSize();
-          results.push({
-            id,
-            rawBase64,
-            width: size.width * pixelRatio,
-            height: size.height * pixelRatio,
-          });
-        }
-        return results;
-      } catch {
-        console.warn('Could not use takeElementScreenshot. Falling back to full-page screenshot and cropping');
-        return this.takePermutationScreenshots();
       }
+      return results;
     } else {
-      // Single full-page screenshot with bounding box metadata for cropping
-
-      const screenshot = await this.fullPageScreenshot();
-      const image = await parsePng(screenshot);
+      const rawBase64 = await this.fullPageScreenshot();
+      const image = await parsePng(rawBase64);
       const permutations = await this.browser.execute(getPermutationSizes);
 
-      return permutations.map((permutation: PermutationInfo) => ({ ...permutation, image, rawBase64: screenshot }));
+      if (permutations.length === 0) {
+        throw new Error('No permutations found on current page.');
+      }
+
+      return permutations.map((permutation: PermutationInfo) => ({
+        id: permutation.id,
+        image,
+        offset: permutation.offset,
+        width: permutation.width,
+        height: permutation.height,
+        rawBase64,
+      }));
     }
   }
 
diff --git a/src/page-objects/types.ts b/src/page-objects/types.ts
index 6840f5a..98cdbb1 100644
--- a/src/page-objects/types.ts
+++ b/src/page-objects/types.ts
@@ -29,18 +29,24 @@ export interface ScreenshotCapturingOptions {
   singleElements?: boolean;
 }
 
-export interface Screenshot extends ElementSize {
+/**
+ * A raw screenshot with base64 data and dimensions. No decoded image, no offset.
+ * Returned when singleElements is true (takeElementScreenshot was used).
+ */
+export interface RawScreenshot extends ElementSize {
   image?: PNG;
-  pixelRatio?: number;
-  /**
-   * The raw base64-encoded PNG from WebDriver, retained for fast byte-equality
-   * comparison in cropAndCompare. When two screenshots have the same rawBase64
-   * and no cropping is needed, expensive decoding is skipped entirely.
-   */
   rawBase64: string;
-  offset?: ElementOffset;
+  pixelRatio?: number;
 }
 
-export interface ScreenshotWithOffset extends Screenshot {
+/**
+ * A decoded screenshot with image data, offset for cropping, and optional rawBase64.
+ * Returned when singleElements is false/absent (full-page screenshot path).
+ */
+export interface ScreenshotWithOffset extends RawScreenshot {
+  image: PNG;
   offset: ElementOffset;
 }
+
+/** Union of both screenshot types. */
+export type Screenshot = RawScreenshot | ScreenshotWithOffset;

From 87eccf5bdbc309ce68993f29e1c79b81183025ee Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Wed, 24 Jun 2026 13:37:19 +0200
Subject: [PATCH 12/23] Type fixes

---
 src/page-objects/index.ts      |  7 +------
 src/page-objects/screenshot.ts | 18 ++++++++++--------
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/src/page-objects/index.ts b/src/page-objects/index.ts
index da26948..216de51 100644
--- a/src/page-objects/index.ts
+++ b/src/page-objects/index.ts
@@ -1,11 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 export { default as BasePageObject } from './base';
-export {
-  default as ScreenshotPageObject,
-  PermutationScreenshot,
-  RawPermutationScreenshot,
-  DecodedPermutationScreenshot,
-} from './screenshot';
+export { default as ScreenshotPageObject, PermutationScreenshot, RawPermutationScreenshot } from './screenshot';
 export { default as EventsSpy } from './events-spy';
 export { Screenshot, RawScreenshot, ScreenshotWithOffset, ElementSize, ElementRect, ElementOffset } from './types';
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 880903a..61f4bc9 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -11,6 +11,7 @@ import { parsePng } from '../image-utils';
 import BasePageObject from './base';
 import { ElementOffset, RawScreenshot, ScreenshotWithOffset, ScreenshotCapturingOptions } from './types';
 import fullPageScreenshot from './full-page-screenshot';
+import { Browser } from 'webdriverio';
 
 /** Permutation screenshot from takeElementScreenshot — no image, no offset. */
 export interface RawPermutationScreenshot extends RawScreenshot {
@@ -18,14 +19,12 @@ export interface RawPermutationScreenshot extends RawScreenshot {
 }
 
 /** Permutation screenshot from full-page fallback — has image and offset for cropping. */
-export interface DecodedPermutationScreenshot extends ScreenshotWithOffset {
+export interface PermutationScreenshot extends ScreenshotWithOffset {
   id: string;
 }
 
-export type PermutationScreenshot = RawPermutationScreenshot | DecodedPermutationScreenshot;
-
 export default class ScreenshotPageObject extends BasePageObject {
-  constructor(browser: WebdriverIO.Browser, public readonly forceScrollAndMerge: boolean = false) {
+  constructor(browser: Browser, public readonly forceScrollAndMerge: boolean = false) {
     super(browser);
   }
 
@@ -79,7 +78,6 @@ export default class ScreenshotPageObject extends BasePageObject {
       return { rawBase64, pixelRatio, height: box.height, width: box.width };
     }
 
-    // Default: full-page or viewport screenshot with decoded image
     const { top, left } = await this.getViewportSize();
     const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
     const image = await parsePng(rawBase64);
@@ -102,8 +100,10 @@ export default class ScreenshotPageObject extends BasePageObject {
 
   // Overloads for capturePermutations
   async capturePermutations(options: { singleElements: true }): Promise<RawPermutationScreenshot[]>;
-  async capturePermutations(options?: { singleElements?: false }): Promise<DecodedPermutationScreenshot[]>;
-  async capturePermutations(options?: { singleElements?: boolean }): Promise<PermutationScreenshot[]> {
+  async capturePermutations(options?: { singleElements?: false }): Promise<PermutationScreenshot[]>;
+  async capturePermutations(options?: {
+    singleElements?: boolean;
+  }): Promise<RawPermutationScreenshot[] | PermutationScreenshot[]> {
     await this.windowScrollTo({ top: 0, left: 0 });
 
     // Adapt viewport height to fit all elements before taking screenshots
@@ -119,7 +119,9 @@ export default class ScreenshotPageObject extends BasePageObject {
     }
   }
 
-  private async takePermutationScreenshots(options?: { singleElements?: boolean }): Promise<PermutationScreenshot[]> {
+  private async takePermutationScreenshots(options?: {
+    singleElements?: boolean;
+  }): Promise<RawPermutationScreenshot[] | PermutationScreenshot[]> {
     if (options?.singleElements) {
       const elements = this.browser.$$('[data-permutation]');
       if ((await elements.length) === 0) {

From cf5e9b98a3857a17c2993071abd1dcfbd101beed Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Wed, 24 Jun 2026 14:52:01 +0200
Subject: [PATCH 13/23] Type fixes

---
 src/image-utils/compare.ts     |  7 ++++---
 src/page-objects/screenshot.ts | 14 ++++++++++----
 src/page-objects/types.ts      |  6 ++++--
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index 15c9c00..b1fb71c 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -50,10 +50,11 @@ export interface CropAndCompareResult {
   diffPixels: number;
 }
 
-async function getScreenshotImage(screenshot: Screenshot) {
-  if (!screenshot.image) {
-    screenshot.image = await parsePng(screenshot.rawBase64);
+async function getScreenshotImage(screenshot: Screenshot): Promise<PNG> {
+  if (isScreenshotWithOffset(screenshot)) {
+    return screenshot.image;
   }
+  screenshot.image = screenshot.image || (await parsePng(screenshot.rawBase64));
   return screenshot.image;
 }
 
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 61f4bc9..fa35bc6 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -91,11 +91,18 @@ export default class ScreenshotPageObject extends BasePageObject {
     return { image, offset, pixelRatio, height: box.height, width: box.width, rawBase64 };
   }
 
-  async captureViewport(): Promise<ScreenshotWithOffset> {
+  async captureViewport(options: { singleElements: true }): Promise<RawScreenshot>;
+  async captureViewport(options?: { singleElements?: false }): Promise<ScreenshotWithOffset>;
+  async captureViewport(options?: { singleElements?: boolean }): Promise<RawScreenshot | ScreenshotWithOffset> {
     const { height, width } = await this.getViewportSize();
     const rawBase64 = await this.browser.takeScreenshot();
+
+    if (options?.singleElements) {
+      return { rawBase64, height, width };
+    }
+
     const image = await parsePng(rawBase64);
-    return { image, offset: { top: 0, left: 0 }, height, width, rawBase64 };
+    return { image, offset: { top: 0, left: 0 }, height, width };
   }
 
   // Overloads for capturePermutations
@@ -121,7 +128,7 @@ export default class ScreenshotPageObject extends BasePageObject {
 
   private async takePermutationScreenshots(options?: {
     singleElements?: boolean;
-  }): Promise<RawPermutationScreenshot[] | PermutationScreenshot[]> {
+  }): Promise<PermutationScreenshot[] | RawPermutationScreenshot[]> {
     if (options?.singleElements) {
       const elements = this.browser.$$('[data-permutation]');
       if ((await elements.length) === 0) {
@@ -159,7 +166,6 @@ export default class ScreenshotPageObject extends BasePageObject {
         offset: permutation.offset,
         width: permutation.width,
         height: permutation.height,
-        rawBase64,
       }));
     }
   }
diff --git a/src/page-objects/types.ts b/src/page-objects/types.ts
index 98cdbb1..760ed72 100644
--- a/src/page-objects/types.ts
+++ b/src/page-objects/types.ts
@@ -40,12 +40,14 @@ export interface RawScreenshot extends ElementSize {
 }
 
 /**
- * A decoded screenshot with image data, offset for cropping, and optional rawBase64.
+ * A decoded screenshot with image data and offset for cropping.
  * Returned when singleElements is false/absent (full-page screenshot path).
+ * This is the original type consumers expect.
  */
-export interface ScreenshotWithOffset extends RawScreenshot {
+export interface ScreenshotWithOffset extends ElementSize {
   image: PNG;
   offset: ElementOffset;
+  pixelRatio?: number;
 }
 
 /** Union of both screenshot types. */

From 206b46232cc36e23adb8cf42899bcc993cf4b07d Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Wed, 24 Jun 2026 21:38:21 +0200
Subject: [PATCH 14/23] Fix and add tests

---
 test/page-object.test.ts | 58 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/test/page-object.test.ts b/test/page-object.test.ts
index ff099bd..03b6cda 100644
--- a/test/page-object.test.ts
+++ b/test/page-object.test.ts
@@ -373,11 +373,16 @@ describe('capturePermutations', () => {
       expect(permutations.map(p => p.id)).toEqual(['perm-1', 'perm-2', 'perm-3']);
 
       permutations.forEach(perm => {
-        expect(perm.rawBase64).toBeDefined();
-        expect(perm.rawBase64.length).toBeGreaterThan(0);
+        expect(perm.image).toBeDefined();
         expect(perm.width).toBeGreaterThan(0);
         expect(perm.height).toBeGreaterThan(0);
+        expect(perm.offset.top).toBeGreaterThanOrEqual(0);
+        expect(perm.offset.left).toBeGreaterThanOrEqual(0);
       });
+
+      // Verify permutations are in order (top to bottom)
+      expect(permutations[0].offset.top).toBeLessThan(permutations[1].offset.top);
+      expect(permutations[1].offset.top).toBeLessThan(permutations[2].offset.top);
     }, './test-permutations.html')
   );
 
@@ -388,3 +393,52 @@ describe('capturePermutations', () => {
     })
   );
 });
+
+describe('singleElements option', () => {
+  test(
+    'captureBySelector with singleElements returns rawBase64 without image',
+    setupTest(async page => {
+      const result = await page.captureBySelector('#text-content', { singleElements: true });
+
+      expect(result.rawBase64).toBeDefined();
+      expect(result.rawBase64.length).toBeGreaterThan(0);
+      expect(result.width).toBeGreaterThan(0);
+      expect(result.height).toBeGreaterThan(0);
+      expect('image' in result).toBe(false);
+      expect('offset' in result).toBe(false);
+    })
+  );
+
+  test(
+    'captureViewport with singleElements returns rawBase64 without image',
+    setupTest(async page => {
+      const result = await page.captureViewport({ singleElements: true });
+
+      expect(result.rawBase64).toBeDefined();
+      expect(result.rawBase64.length).toBeGreaterThan(0);
+      expect(result.width).toBeGreaterThan(0);
+      expect(result.height).toBeGreaterThan(0);
+      expect('image' in result).toBe(false);
+      expect('offset' in result).toBe(false);
+    })
+  );
+
+  test(
+    'capturePermutations with singleElements returns rawBase64 without image',
+    setupTest(async page => {
+      const permutations = await page.capturePermutations({ singleElements: true });
+
+      expect(permutations).toHaveLength(3);
+      expect(permutations.map(p => p.id)).toEqual(['perm-1', 'perm-2', 'perm-3']);
+
+      permutations.forEach(perm => {
+        expect(perm.rawBase64).toBeDefined();
+        expect(perm.rawBase64.length).toBeGreaterThan(0);
+        expect(perm.width).toBeGreaterThan(0);
+        expect(perm.height).toBeGreaterThan(0);
+        expect('image' in perm).toBe(false);
+        expect('offset' in perm).toBe(false);
+      });
+    }, './test-permutations.html')
+  );
+});

From 3ab29956f7d8c7a606925571dd8d8554a396ee65 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Wed, 24 Jun 2026 23:08:55 +0200
Subject: [PATCH 15/23] Fixes

---
 src/image-utils/compare.ts     | 8 +++++---
 src/page-objects/screenshot.ts | 4 ++--
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index b1fb71c..df5b319 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -54,6 +54,8 @@ async function getScreenshotImage(screenshot: Screenshot): Promise<PNG> {
   if (isScreenshotWithOffset(screenshot)) {
     return screenshot.image;
   }
+  // Cache the parsed image in the screenshot object so that it does not need to be parsed anymore
+  // when cropping it again for a different offset
   screenshot.image = screenshot.image || (await parsePng(screenshot.rawBase64));
   return screenshot.image;
 }
@@ -86,10 +88,10 @@ export async function cropAndCompare(
   const size = normalizeSize(firstScreenshot, secondScreenshot);
   const scaledSize = scaleSize(size, pixelRatio);
 
-  const firstImage = await cropIfNeeded(firstScreenshot, scaledSize);
-  const secondImage = await cropIfNeeded(secondScreenshot, scaledSize);
+  const firstImage = await cropIfNeeded(firstScreenshot, size);
+  const secondImage = await cropIfNeeded(secondScreenshot, size);
 
-  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, size);
+  const { diffImage, diffPixels } = compareImages(firstImage, secondImage, scaledSize);
 
   // Skip packPng when no cropping was needed and rawBase64 is available
   const [firstPacked, secondPacked, diffPacked] = await Promise.all([
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index fa35bc6..b9fe09c 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -130,7 +130,7 @@ export default class ScreenshotPageObject extends BasePageObject {
     singleElements?: boolean;
   }): Promise<PermutationScreenshot[] | RawPermutationScreenshot[]> {
     if (options?.singleElements) {
-      const elements = this.browser.$$('[data-permutation]');
+      const elements = await this.browser.$$('[data-permutation]');
       if ((await elements.length) === 0) {
         throw new Error('No permutations found on current page.');
       }
@@ -141,7 +141,7 @@ export default class ScreenshotPageObject extends BasePageObject {
       const results: RawPermutationScreenshot[] = [];
       for (const element of elements) {
         const id = (await element.getAttribute('data-permutation')) || '';
-        const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
+        const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
         const size = await element.getSize();
         results.push({
           id,

From a82849514371cf618d13fe8b1375c05dafbe6405 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Wed, 24 Jun 2026 23:39:31 +0200
Subject: [PATCH 16/23] Fix getting pixelRatio

---
 src/page-objects/screenshot.ts | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index b9fe09c..4798e0e 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -130,18 +130,16 @@ export default class ScreenshotPageObject extends BasePageObject {
     singleElements?: boolean;
   }): Promise<PermutationScreenshot[] | RawPermutationScreenshot[]> {
     if (options?.singleElements) {
-      const elements = await this.browser.$$('[data-permutation]');
+      const elements = this.browser.$$('[data-permutation]');
       if ((await elements.length) === 0) {
         throw new Error('No permutations found on current page.');
       }
 
-      const pixelRatio = await this.browser.execute(function () {
-        return window.devicePixelRatio || 1;
-      });
+      const { pixelRatio } = await this.getViewportSize();
       const results: RawPermutationScreenshot[] = [];
       for (const element of elements) {
         const id = (await element.getAttribute('data-permutation')) || '';
-        const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
+        const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
         const size = await element.getSize();
         results.push({
           id,

From 1a4f41e0e06ea931c55c7979c4b992c6cfdbaac4 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Wed, 24 Jun 2026 23:46:02 +0200
Subject: [PATCH 17/23] Fix

---
 src/page-objects/screenshot.ts | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index 4798e0e..af50be5 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -130,8 +130,8 @@ export default class ScreenshotPageObject extends BasePageObject {
     singleElements?: boolean;
   }): Promise<PermutationScreenshot[] | RawPermutationScreenshot[]> {
     if (options?.singleElements) {
-      const elements = this.browser.$$('[data-permutation]');
-      if ((await elements.length) === 0) {
+      const elements = await this.browser.$$('[data-permutation]').map(el => el);
+      if (elements.length === 0) {
         throw new Error('No permutations found on current page.');
       }
 
@@ -139,7 +139,7 @@ export default class ScreenshotPageObject extends BasePageObject {
       const results: RawPermutationScreenshot[] = [];
       for (const element of elements) {
         const id = (await element.getAttribute('data-permutation')) || '';
-        const rawBase64 = await this.browser.takeElementScreenshot(element.elementId);
+        const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
         const size = await element.getSize();
         results.push({
           id,

From a60a9745e7615d4aaa12a112ed818f42e124dad9 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Thu, 25 Jun 2026 00:42:16 +0200
Subject: [PATCH 18/23] Fixes

---
 src/image-utils/compare.ts | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/image-utils/compare.ts b/src/image-utils/compare.ts
index df5b319..63c6058 100644
--- a/src/image-utils/compare.ts
+++ b/src/image-utils/compare.ts
@@ -3,13 +3,15 @@
 import { PNG } from 'pngjs';
 import pixelmatch from 'pixelmatch';
 import { packPng, cropImage, parsePng } from './utils';
-import { ElementRect, ElementSize, Screenshot, ScreenshotWithOffset } from '../page-objects/types';
+import { ElementRect, ElementSize, RawScreenshot, ScreenshotWithOffset } from '../page-objects/types';
 
 interface Size {
   width: number;
   height: number;
 }
 
+type Screenshot = ScreenshotWithOffset | RawScreenshot;
+
 export function compareImages(firstImage: PNG, secondImage: PNG, { width, height }: ElementSize) {
   // This prevents an error thrown from pixelmatch when comparing 0-sized images.
   if (width === 0 || height === 0) {
@@ -50,18 +52,15 @@ export interface CropAndCompareResult {
   diffPixels: number;
 }
 
-async function getScreenshotImage(screenshot: Screenshot): Promise<PNG> {
+async function getDecodedImage(screenshot: Screenshot): Promise<PNG> {
   if (isScreenshotWithOffset(screenshot)) {
     return screenshot.image;
   }
-  // Cache the parsed image in the screenshot object so that it does not need to be parsed anymore
-  // when cropping it again for a different offset
-  screenshot.image = screenshot.image || (await parsePng(screenshot.rawBase64));
-  return screenshot.image;
+  return parsePng(screenshot.rawBase64);
 }
 
 async function cropIfNeeded(screenshot: Screenshot, size: Size) {
-  const image = await getScreenshotImage(screenshot);
+  const image = await getDecodedImage(screenshot);
   if (isScreenshotWithOffset(screenshot)) {
     return cropImage(image, buildCropRect(screenshot as ScreenshotWithOffset, size), screenshot.pixelRatio);
   } else {

From 5703781129e0da4e77bb8ddbc8ac8236704da6ad Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Thu, 25 Jun 2026 17:13:14 +0200
Subject: [PATCH 19/23] Refactor tests

---
 src/page-objects/index.ts               |   6 +-
 src/page-objects/raw-screenshot.ts      |  68 +++++++++
 src/page-objects/screenshot-base.ts     |  61 ++++++++
 src/page-objects/screenshot.ts          | 185 ++++--------------------
 src/page-objects/types.ts               |  26 ++--
 test/page-object.test.ts                | 114 ---------------
 test/raw-screenshot-page-object.test.ts |  70 +++++++++
 test/screenshot-page-object.test.ts     |  79 ++++++++++
 8 files changed, 319 insertions(+), 290 deletions(-)
 create mode 100644 src/page-objects/raw-screenshot.ts
 create mode 100644 src/page-objects/screenshot-base.ts
 create mode 100644 test/raw-screenshot-page-object.test.ts
 create mode 100644 test/screenshot-page-object.test.ts

diff --git a/src/page-objects/index.ts b/src/page-objects/index.ts
index 216de51..2af8870 100644
--- a/src/page-objects/index.ts
+++ b/src/page-objects/index.ts
@@ -1,6 +1,8 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 export { default as BasePageObject } from './base';
-export { default as ScreenshotPageObject, PermutationScreenshot, RawPermutationScreenshot } from './screenshot';
+export { default as ScreenshotBasePageObject } from './screenshot-base';
+export { default as ScreenshotPageObject, PermutationScreenshot } from './screenshot';
+export { default as RawScreenshotPageObject, RawPermutationScreenshot } from './raw-screenshot';
 export { default as EventsSpy } from './events-spy';
-export { Screenshot, RawScreenshot, ScreenshotWithOffset, ElementSize, ElementRect, ElementOffset } from './types';
+export { ScreenshotWithOffset, RawScreenshot, ElementSize, ElementRect, ElementOffset } from './types';
diff --git a/src/page-objects/raw-screenshot.ts b/src/page-objects/raw-screenshot.ts
new file mode 100644
index 0000000..7c92d62
--- /dev/null
+++ b/src/page-objects/raw-screenshot.ts
@@ -0,0 +1,68 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import ScreenshotBasePageObject from './screenshot-base';
+import { RawScreenshot } from './types';
+
+/**
+ * Raw permutation screenshot captured via takeElementScreenshot.
+ * No decoded image, no offset — just the raw base64 PNG per element.
+ */
+export interface RawPermutationScreenshot extends RawScreenshot {
+  id: string;
+}
+
+/**
+ * A page object that captures screenshots using takeElementScreenshot.
+ * Returns raw base64 PNGs without decoding, cropping, or re-encoding —
+ * significantly faster when pixel-level comparison can be done on raw bytes.
+ */
+export default class RawScreenshotPageObject extends ScreenshotBasePageObject {
+  async captureBySelector(selector: string): Promise<RawScreenshot> {
+    await this.waitForVisible(selector);
+    const { pixelRatio } = await this.getViewportSize();
+    const box = await this.getBoundingBox(selector);
+
+    const originalWindowSize = await this.fitWindowHeightToContent();
+    const element = this.browser.$(selector);
+    const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
+    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+
+    return { rawBase64, pixelRatio, height: box.height, width: box.width };
+  }
+
+  async captureViewport(): Promise<RawScreenshot> {
+    const { height, width } = await this.getViewportSize();
+    const rawBase64 = await this.browser.takeScreenshot();
+    return { rawBase64, height, width };
+  }
+
+  async capturePermutations(): Promise<RawPermutationScreenshot[]> {
+    await this.windowScrollTo({ top: 0, left: 0 });
+
+    // Adapt viewport height to fit all elements before taking screenshots
+    const originalWindowSize = await this.fitWindowHeightToContent();
+
+    const elements = await this.browser.$$('[data-permutation]').map(el => el);
+    if (elements.length === 0) {
+      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+      throw new Error('No permutations found on current page.');
+    }
+
+    const { pixelRatio } = await this.getViewportSize();
+    const results: RawPermutationScreenshot[] = [];
+    for (const element of elements) {
+      const id = (await element.getAttribute('data-permutation')) || '';
+      const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
+      const size = await element.getSize();
+      results.push({
+        id,
+        rawBase64,
+        width: size.width * pixelRatio,
+        height: size.height * pixelRatio,
+      });
+    }
+
+    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
+    return results;
+  }
+}
diff --git a/src/page-objects/screenshot-base.ts b/src/page-objects/screenshot-base.ts
new file mode 100644
index 0000000..dce2e6a
--- /dev/null
+++ b/src/page-objects/screenshot-base.ts
@@ -0,0 +1,61 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import { ScrollAction, scrollAction, getPageDimensions } from '../browser-scripts';
+import BasePageObject from './base';
+import fullPageScreenshot from './full-page-screenshot';
+
+/**
+ * Base class for screenshot page objects. Provides scroll helpers,
+ * full-page screenshot support, and window size management.
+ */
+export default class ScreenshotBasePageObject extends BasePageObject {
+  constructor(browser: WebdriverIO.Browser, public readonly forceScrollAndMerge: boolean = false) {
+    super(browser);
+  }
+
+  async focusNextElement() {
+    return this.keys('Tab');
+  }
+
+  async scrollToBottom(selector: string) {
+    const action: ScrollAction = 'scrollToBottom';
+    await this.browser.execute(scrollAction, { action, selector });
+  }
+
+  async scrollToRight(selector: string) {
+    const action: ScrollAction = 'scrollToRight';
+    await this.browser.execute(scrollAction, { action, selector });
+  }
+
+  async fullPageScreenshot() {
+    // preserve scroll position in order to avoid side effects after screenshot taking
+    const scrollPosition = await this.getWindowScroll();
+    // Wait for the page to settle before taking a screenshot
+    await this.waitForJsTimers();
+    const screenshot = await fullPageScreenshot(this.browser, this.forceScrollAndMerge);
+    // restore scroll position
+    await this.windowScrollTo(scrollPosition);
+    return screenshot;
+  }
+
+  protected async fitWindowHeightToContent(): Promise<{ width: number; height: number }> {
+    const originalWindowSize = await this.browser.getWindowSize();
+    const { viewportHeight, pageHeight } = await this.browser.execute(getPageDimensions);
+    const windowUIHeight = originalWindowSize.height - viewportHeight;
+    await this.safeSetWindowSize(originalWindowSize.width, pageHeight + windowUIHeight);
+    return originalWindowSize;
+  }
+
+  /* istanbul ignore next -- setWindowSize is unsupported on some mobile browsers, not testable in CI */
+  protected async safeSetWindowSize(width: number, height: number): Promise<void> {
+    try {
+      await this.browser.setWindowSize(width, height);
+    } catch (error) {
+      if (error instanceof Error && error.message.includes('Method has not yet been implemented')) {
+        console.log('setWindowSize is not supported on this device');
+      } else {
+        throw error;
+      }
+    }
+  }
+}
diff --git a/src/page-objects/screenshot.ts b/src/page-objects/screenshot.ts
index af50be5..472c43d 100644
--- a/src/page-objects/screenshot.ts
+++ b/src/page-objects/screenshot.ts
@@ -1,191 +1,62 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import {
-  ScrollAction,
-  scrollAction,
-  getPermutationSizes,
-  getPageDimensions,
-  PermutationInfo,
-} from '../browser-scripts';
+import { getPermutationSizes, PermutationInfo } from '../browser-scripts';
 import { parsePng } from '../image-utils';
-import BasePageObject from './base';
-import { ElementOffset, RawScreenshot, ScreenshotWithOffset, ScreenshotCapturingOptions } from './types';
-import fullPageScreenshot from './full-page-screenshot';
-import { Browser } from 'webdriverio';
+import { ElementOffset, ScreenshotCapturingOptions, ScreenshotWithOffset } from './types';
+import ScreenshotBasePageObject from './screenshot-base';
 
-/** Permutation screenshot from takeElementScreenshot — no image, no offset. */
-export interface RawPermutationScreenshot extends RawScreenshot {
-  id: string;
-}
-
-/** Permutation screenshot from full-page fallback — has image and offset for cropping. */
 export interface PermutationScreenshot extends ScreenshotWithOffset {
   id: string;
 }
 
-export default class ScreenshotPageObject extends BasePageObject {
-  constructor(browser: Browser, public readonly forceScrollAndMerge: boolean = false) {
-    super(browser);
-  }
-
-  async focusNextElement() {
-    return this.keys('Tab');
-  }
-
-  async scrollToBottom(selector: string) {
-    const action: ScrollAction = 'scrollToBottom';
-    await this.browser.execute(scrollAction, { action, selector });
-  }
-
-  async scrollToRight(selector: string) {
-    const action: ScrollAction = 'scrollToRight';
-    await this.browser.execute(scrollAction, { action, selector });
-  }
-
-  async fullPageScreenshot() {
-    // preserve scroll position in order to avoid side effects after screenshot taking
-    const scrollPosition = await this.getWindowScroll();
-    // Wait for the page to settle before taking a screenshot
-    await this.waitForJsTimers();
-    const screenshot = await fullPageScreenshot(this.browser, this.forceScrollAndMerge);
-    // restore scroll position
-    await this.windowScrollTo(scrollPosition);
-    return screenshot;
-  }
-
-  // Overloads for captureBySelector
-  async captureBySelector(
-    selector: string,
-    options: ScreenshotCapturingOptions & { singleElements: true }
-  ): Promise<RawScreenshot>;
-  async captureBySelector(
-    selector: string,
-    options?: ScreenshotCapturingOptions & { singleElements?: false }
-  ): Promise<ScreenshotWithOffset>;
-  async captureBySelector(
-    selector: string,
-    options: ScreenshotCapturingOptions = {}
-  ): Promise<RawScreenshot | ScreenshotWithOffset> {
+export default class ScreenshotPageObject extends ScreenshotBasePageObject {
+  async captureBySelector(selector: string, options: ScreenshotCapturingOptions = {}): Promise<ScreenshotWithOffset> {
     await this.waitForVisible(selector);
-    const { pixelRatio } = await this.getViewportSize();
+    const { pixelRatio, top, left } = await this.getViewportSize();
     const box = await this.getBoundingBox(selector);
-
-    if (options.singleElements) {
-      const originalWindowSize = await this.fitWindowHeightToContent();
-      const element = this.browser.$(selector);
-      const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
-      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
-      return { rawBase64, pixelRatio, height: box.height, width: box.width };
-    }
-
-    const { top, left } = await this.getViewportSize();
-    const rawBase64 = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
-    const image = await parsePng(rawBase64);
+    const screenshot = options.viewportOnly ? await this.browser.takeScreenshot() : await this.fullPageScreenshot();
+    const image = await parsePng(screenshot);
 
     const offset: ElementOffset = { top: box.top, left: box.left };
     if (!options.viewportOnly) {
+      // Correct potential scrolling offsets when using a full page screenshot
       offset.top += top;
       offset.left += left;
     }
 
-    return { image, offset, pixelRatio, height: box.height, width: box.width, rawBase64 };
+    return { image, offset, pixelRatio, height: box.height, width: box.width };
   }
 
-  async captureViewport(options: { singleElements: true }): Promise<RawScreenshot>;
-  async captureViewport(options?: { singleElements?: false }): Promise<ScreenshotWithOffset>;
-  async captureViewport(options?: { singleElements?: boolean }): Promise<RawScreenshot | ScreenshotWithOffset> {
+  async captureViewport(): Promise<ScreenshotWithOffset> {
     const { height, width } = await this.getViewportSize();
-    const rawBase64 = await this.browser.takeScreenshot();
 
-    if (options?.singleElements) {
-      return { rawBase64, height, width };
-    }
+    const offset: ElementOffset = {
+      top: 0,
+      left: 0,
+    };
 
-    const image = await parsePng(rawBase64);
-    return { image, offset: { top: 0, left: 0 }, height, width };
+    const screenshot = await this.browser.takeScreenshot();
+    const image = await parsePng(screenshot);
+    return { image, offset, height, width };
   }
 
-  // Overloads for capturePermutations
-  async capturePermutations(options: { singleElements: true }): Promise<RawPermutationScreenshot[]>;
-  async capturePermutations(options?: { singleElements?: false }): Promise<PermutationScreenshot[]>;
-  async capturePermutations(options?: {
-    singleElements?: boolean;
-  }): Promise<RawPermutationScreenshot[] | PermutationScreenshot[]> {
+  async capturePermutations(): Promise<PermutationScreenshot[]> {
     await this.windowScrollTo({ top: 0, left: 0 });
 
-    // Adapt viewport height to fit all elements before taking screenshots
+    // Adapt viewport height to fit all elements before taking a screenshot
     const originalWindowSize = await this.fitWindowHeightToContent();
 
-    try {
-      const results = await this.takePermutationScreenshots(options);
-      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
-      return results;
-    } catch (error) {
-      await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
-      throw error;
-    }
-  }
-
-  private async takePermutationScreenshots(options?: {
-    singleElements?: boolean;
-  }): Promise<PermutationScreenshot[] | RawPermutationScreenshot[]> {
-    if (options?.singleElements) {
-      const elements = await this.browser.$$('[data-permutation]').map(el => el);
-      if (elements.length === 0) {
-        throw new Error('No permutations found on current page.');
-      }
-
-      const { pixelRatio } = await this.getViewportSize();
-      const results: RawPermutationScreenshot[] = [];
-      for (const element of elements) {
-        const id = (await element.getAttribute('data-permutation')) || '';
-        const rawBase64 = await this.browser.takeElementScreenshot(await element.elementId);
-        const size = await element.getSize();
-        results.push({
-          id,
-          rawBase64,
-          width: size.width * pixelRatio,
-          height: size.height * pixelRatio,
-        });
-      }
-      return results;
-    } else {
-      const rawBase64 = await this.fullPageScreenshot();
-      const image = await parsePng(rawBase64);
-      const permutations = await this.browser.execute(getPermutationSizes);
+    const screenshot = await this.fullPageScreenshot();
+    const image = await parsePng(screenshot);
+    const permutations = await this.browser.execute(getPermutationSizes);
 
-      if (permutations.length === 0) {
-        throw new Error('No permutations found on current page.');
-      }
+    // Restore window size after taking the screenshot
+    await this.safeSetWindowSize(originalWindowSize.width, originalWindowSize.height);
 
-      return permutations.map((permutation: PermutationInfo) => ({
-        id: permutation.id,
-        image,
-        offset: permutation.offset,
-        width: permutation.width,
-        height: permutation.height,
-      }));
+    if (permutations.length === 0) {
+      throw new Error('No permutations found on current page.');
     }
-  }
 
-  private async fitWindowHeightToContent(): Promise<{ width: number; height: number }> {
-    const originalWindowSize = await this.browser.getWindowSize();
-    const { viewportHeight, pageHeight } = await this.browser.execute(getPageDimensions);
-    const windowUIHeight = originalWindowSize.height - viewportHeight;
-    await this.safeSetWindowSize(originalWindowSize.width, pageHeight + windowUIHeight);
-    return originalWindowSize;
-  }
-
-  /* istanbul ignore next -- setWindowSize is unsupported on some mobile browsers, not testable in CI */
-  private async safeSetWindowSize(width: number, height: number): Promise<void> {
-    try {
-      await this.browser.setWindowSize(width, height);
-    } catch (error) {
-      if (error instanceof Error && error.message.includes('Method has not yet been implemented')) {
-        console.log('setWindowSize is not supported on this device');
-      } else {
-        throw error;
-      }
-    }
+    return permutations.map((permutation: PermutationInfo) => ({ ...permutation, image }));
   }
 }
diff --git a/src/page-objects/types.ts b/src/page-objects/types.ts
index 760ed72..4444c10 100644
--- a/src/page-objects/types.ts
+++ b/src/page-objects/types.ts
@@ -26,29 +26,21 @@ export interface ViewportSize extends ElementOffset, ElementSize {
 
 export interface ScreenshotCapturingOptions {
   viewportOnly?: boolean;
-  singleElements?: boolean;
 }
 
-/**
- * A raw screenshot with base64 data and dimensions. No decoded image, no offset.
- * Returned when singleElements is true (takeElementScreenshot was used).
- */
-export interface RawScreenshot extends ElementSize {
-  image?: PNG;
-  rawBase64: string;
+export interface ScreenshotWithOffset extends ElementSize {
+  image: PNG;
+  offset: ElementOffset;
   pixelRatio?: number;
+  /** Optional raw base64 PNG for fast byte-equality comparison. */
+  rawBase64?: string;
 }
 
 /**
- * A decoded screenshot with image data and offset for cropping.
- * Returned when singleElements is false/absent (full-page screenshot path).
- * This is the original type consumers expect.
+ * A raw screenshot with base64 data and dimensions. No decoded image, no offset.
+ * Returned by RawScreenshotPageObject which uses takeElementScreenshot.
  */
-export interface ScreenshotWithOffset extends ElementSize {
-  image: PNG;
-  offset: ElementOffset;
+export interface RawScreenshot extends ElementSize {
+  rawBase64: string;
   pixelRatio?: number;
 }
-
-/** Union of both screenshot types. */
-export type Screenshot = RawScreenshot | ScreenshotWithOffset;
diff --git a/test/page-object.test.ts b/test/page-object.test.ts
index 03b6cda..8ef9b27 100644
--- a/test/page-object.test.ts
+++ b/test/page-object.test.ts
@@ -48,15 +48,6 @@ test(
   })
 );
 
-test(
-  'focusNextElement',
-  setupTest(async page => {
-    await page.click('#input-1');
-    await page.focusNextElement();
-    expect(await page.isFocused('#input-2')).toBe(true);
-  })
-);
-
 test(
   'isSelected',
   setupTest(async page => {
@@ -269,31 +260,6 @@ test(
   })
 );
 
-test(
-  'scrollToRight',
-  setupTest(async page => {
-    const width = 400;
-    const overscroll = (width * 20) / 100; // The container has 120% width
-    await page.setWindowSize({ width, height: 300 });
-    await page.scrollToRight('#scrollable-container');
-    expect(await page.getElementScroll('#scrollable-container')).toEqual({ top: 0, left: overscroll });
-  })
-);
-
-test(
-  'scrollToBottom',
-  setupTest(async page => {
-    const testElement = '#scrollable-container';
-    const { height } = await page.getBoundingBox(testElement);
-    const scrollHeight = (await page.getElementProperty(testElement, 'scrollHeight')) as number;
-
-    await page.scrollToBottom(testElement);
-
-    const { top } = await page.getElementScroll(testElement);
-    expect(top).toBeGreaterThanOrEqual(scrollHeight - height);
-  })
-);
-
 test(
   'getBoundingBox',
   setupTest(async page => {
@@ -362,83 +328,3 @@ test(
     });
   })
 );
-
-describe('capturePermutations', () => {
-  test(
-    'captures all permutations with correct ids and dimensions',
-    setupTest(async page => {
-      const permutations = await page.capturePermutations();
-
-      expect(permutations).toHaveLength(3);
-      expect(permutations.map(p => p.id)).toEqual(['perm-1', 'perm-2', 'perm-3']);
-
-      permutations.forEach(perm => {
-        expect(perm.image).toBeDefined();
-        expect(perm.width).toBeGreaterThan(0);
-        expect(perm.height).toBeGreaterThan(0);
-        expect(perm.offset.top).toBeGreaterThanOrEqual(0);
-        expect(perm.offset.left).toBeGreaterThanOrEqual(0);
-      });
-
-      // Verify permutations are in order (top to bottom)
-      expect(permutations[0].offset.top).toBeLessThan(permutations[1].offset.top);
-      expect(permutations[1].offset.top).toBeLessThan(permutations[2].offset.top);
-    }, './test-permutations.html')
-  );
-
-  test(
-    'throws error when no permutations found',
-    setupTest(async page => {
-      await expect(page.capturePermutations()).rejects.toThrowError('No permutations found on current page.');
-    })
-  );
-});
-
-describe('singleElements option', () => {
-  test(
-    'captureBySelector with singleElements returns rawBase64 without image',
-    setupTest(async page => {
-      const result = await page.captureBySelector('#text-content', { singleElements: true });
-
-      expect(result.rawBase64).toBeDefined();
-      expect(result.rawBase64.length).toBeGreaterThan(0);
-      expect(result.width).toBeGreaterThan(0);
-      expect(result.height).toBeGreaterThan(0);
-      expect('image' in result).toBe(false);
-      expect('offset' in result).toBe(false);
-    })
-  );
-
-  test(
-    'captureViewport with singleElements returns rawBase64 without image',
-    setupTest(async page => {
-      const result = await page.captureViewport({ singleElements: true });
-
-      expect(result.rawBase64).toBeDefined();
-      expect(result.rawBase64.length).toBeGreaterThan(0);
-      expect(result.width).toBeGreaterThan(0);
-      expect(result.height).toBeGreaterThan(0);
-      expect('image' in result).toBe(false);
-      expect('offset' in result).toBe(false);
-    })
-  );
-
-  test(
-    'capturePermutations with singleElements returns rawBase64 without image',
-    setupTest(async page => {
-      const permutations = await page.capturePermutations({ singleElements: true });
-
-      expect(permutations).toHaveLength(3);
-      expect(permutations.map(p => p.id)).toEqual(['perm-1', 'perm-2', 'perm-3']);
-
-      permutations.forEach(perm => {
-        expect(perm.rawBase64).toBeDefined();
-        expect(perm.rawBase64.length).toBeGreaterThan(0);
-        expect(perm.width).toBeGreaterThan(0);
-        expect(perm.height).toBeGreaterThan(0);
-        expect('image' in perm).toBe(false);
-        expect('offset' in perm).toBe(false);
-      });
-    }, './test-permutations.html')
-  );
-});
diff --git a/test/raw-screenshot-page-object.test.ts b/test/raw-screenshot-page-object.test.ts
new file mode 100644
index 0000000..2ff4781
--- /dev/null
+++ b/test/raw-screenshot-page-object.test.ts
@@ -0,0 +1,70 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import { test, expect, describe } from 'vitest';
+import { RawScreenshotPageObject } from '../src/page-objects';
+import useBrowser from '../src/use-browser';
+import './utils/setup-local-driver';
+
+type TestFn = (page: RawScreenshotPageObject) => Promise<void>;
+function setupTest(testFn: TestFn, url = './test-page-object.html') {
+  return useBrowser(async browser => {
+    await browser.url(url);
+    await testFn(new RawScreenshotPageObject(browser));
+  });
+}
+
+describe('RawScreenshotPageObject', () => {
+  test(
+    'captureBySelector returns rawBase64 without image or offset',
+    setupTest(async page => {
+      const result = await page.captureBySelector('#text-content');
+
+      expect(result.rawBase64).toBeDefined();
+      expect(result.rawBase64.length).toBeGreaterThan(0);
+      expect(result.width).toBeGreaterThan(0);
+      expect(result.height).toBeGreaterThan(0);
+      expect('image' in result).toBe(false);
+      expect('offset' in result).toBe(false);
+    })
+  );
+
+  test(
+    'captureViewport returns rawBase64 without image or offset',
+    setupTest(async page => {
+      const result = await page.captureViewport();
+
+      expect(result.rawBase64).toBeDefined();
+      expect(result.rawBase64.length).toBeGreaterThan(0);
+      expect(result.width).toBeGreaterThan(0);
+      expect(result.height).toBeGreaterThan(0);
+      expect('image' in result).toBe(false);
+      expect('offset' in result).toBe(false);
+    })
+  );
+
+  test(
+    'capturePermutations returns rawBase64 without image or offset',
+    setupTest(async page => {
+      const permutations = await page.capturePermutations();
+
+      expect(permutations).toHaveLength(3);
+      expect(permutations.map(p => p.id)).toEqual(['perm-1', 'perm-2', 'perm-3']);
+
+      permutations.forEach(perm => {
+        expect(perm.rawBase64).toBeDefined();
+        expect(perm.rawBase64.length).toBeGreaterThan(0);
+        expect(perm.width).toBeGreaterThan(0);
+        expect(perm.height).toBeGreaterThan(0);
+        expect('image' in perm).toBe(false);
+        expect('offset' in perm).toBe(false);
+      });
+    }, './test-permutations.html')
+  );
+
+  test(
+    'capturePermutations throws when no permutations found',
+    setupTest(async page => {
+      await expect(page.capturePermutations()).rejects.toThrowError('No permutations found on current page.');
+    })
+  );
+});
diff --git a/test/screenshot-page-object.test.ts b/test/screenshot-page-object.test.ts
new file mode 100644
index 0000000..6ee4d67
--- /dev/null
+++ b/test/screenshot-page-object.test.ts
@@ -0,0 +1,79 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import { test, expect, describe } from 'vitest';
+import { ScreenshotPageObject } from '../src/page-objects';
+import useBrowser from '../src/use-browser';
+import './utils/setup-local-driver';
+
+type TestFn = (page: ScreenshotPageObject) => Promise<void>;
+function setupTest(testFn: TestFn, url = './test-page-object.html') {
+  return useBrowser(async browser => {
+    await browser.url(url);
+    await testFn(new ScreenshotPageObject(browser));
+  });
+}
+
+test(
+  'focusNextElement',
+  setupTest(async page => {
+    await page.click('#input-1');
+    await page.focusNextElement();
+    expect(await page.isFocused('#input-2')).toBe(true);
+  })
+);
+
+test(
+  'scrollToRight',
+  setupTest(async page => {
+    const width = 400;
+    const overscroll = (width * 20) / 100; // The container has 120% width
+    await page.setWindowSize({ width, height: 300 });
+    await page.scrollToRight('#scrollable-container');
+    expect(await page.getElementScroll('#scrollable-container')).toEqual({ top: 0, left: overscroll });
+  })
+);
+
+test(
+  'scrollToBottom',
+  setupTest(async page => {
+    const testElement = '#scrollable-container';
+    const { height } = await page.getBoundingBox(testElement);
+    const scrollHeight = (await page.getElementProperty(testElement, 'scrollHeight')) as number;
+
+    await page.scrollToBottom(testElement);
+
+    const { top } = await page.getElementScroll(testElement);
+    expect(top).toBeGreaterThanOrEqual(scrollHeight - height);
+  })
+);
+
+describe('capturePermutations', () => {
+  test(
+    'captures all permutations with correct ids and dimensions',
+    setupTest(async page => {
+      const permutations = await page.capturePermutations();
+
+      expect(permutations).toHaveLength(3);
+      expect(permutations.map(p => p.id)).toEqual(['perm-1', 'perm-2', 'perm-3']);
+
+      permutations.forEach(perm => {
+        expect(perm.image).toBeDefined();
+        expect(perm.width).toBeGreaterThan(0);
+        expect(perm.height).toBeGreaterThan(0);
+        expect(perm.offset.top).toBeGreaterThanOrEqual(0);
+        expect(perm.offset.left).toBeGreaterThanOrEqual(0);
+      });
+
+      // Verify permutations are in order (top to bottom)
+      expect(permutations[0].offset.top).toBeLessThan(permutations[1].offset.top);
+      expect(permutations[1].offset.top).toBeLessThan(permutations[2].offset.top);
+    }, './test-permutations.html')
+  );
+
+  test(
+    'throws error when no permutations found',
+    setupTest(async page => {
+      await expect(page.capturePermutations()).rejects.toThrowError('No permutations found on current page.');
+    })
+  );
+});

From cd40a260a8031a5bc23f700187faf2014d632603 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Thu, 25 Jun 2026 17:22:40 +0200
Subject: [PATCH 20/23] Add coverage

---
 test/compare-images.test.ts | 90 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 88 insertions(+), 2 deletions(-)

diff --git a/test/compare-images.test.ts b/test/compare-images.test.ts
index b6091f0..85b06f7 100644
--- a/test/compare-images.test.ts
+++ b/test/compare-images.test.ts
@@ -1,10 +1,10 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import { test, expect } from 'vitest';
+import { test, expect, describe } from 'vitest';
 import fs from 'fs';
 import { PNG } from 'pngjs';
 import useBrowser from '../src/use-browser';
-import { ScreenshotPageObject, ScreenshotWithOffset } from '../src/page-objects';
+import { ScreenshotPageObject, ScreenshotWithOffset, RawScreenshot } from '../src/page-objects';
 import { cropAndCompare, parsePng } from '../src/image-utils';
 import './utils/setup-local-driver';
 
@@ -258,3 +258,89 @@ test('returns isEqual=false when comparing images with 0-size', async () => {
     })
   );
 });
+
+describe('cropAndCompare with RawScreenshot', () => {
+  test('returns isEqual=true when two identical RawScreenshots are compared', async () => {
+    const rawBase64 = fs.readFileSync(__dirname + '/fixtures/red.png', 'base64');
+    const screenshot: RawScreenshot = { rawBase64, width: 200, height: 100 };
+
+    const result = await cropAndCompare(screenshot, screenshot);
+
+    expect(result.isEqual).toBe(true);
+    expect(result.diffPixels).toBe(0);
+    expect(result.diffImage).toBeNull();
+    expect(result.firstImage).toEqual(Buffer.from(rawBase64, 'base64'));
+    expect(result.secondImage).toEqual(Buffer.from(rawBase64, 'base64'));
+  });
+
+  test('detects differences between two different RawScreenshots', async () => {
+    const redBase64 = fs.readFileSync(__dirname + '/fixtures/red.png', 'base64');
+    const blueBase64 = fs.readFileSync(__dirname + '/fixtures/blue.png', 'base64');
+
+    const first: RawScreenshot = { rawBase64: redBase64, width: 200, height: 100 };
+    const second: RawScreenshot = { rawBase64: blueBase64, width: 200, height: 100 };
+
+    const result = await cropAndCompare(first, second);
+
+    expect(result.isEqual).toBe(false);
+    expect(result.diffPixels).toBeGreaterThan(0);
+    expect(result.firstImage).toBeInstanceOf(Buffer);
+    expect(result.secondImage).toBeInstanceOf(Buffer);
+  });
+
+  test('compares a RawScreenshot against a ScreenshotWithOffset', async () => {
+    const rawBase64 = fs.readFileSync(__dirname + '/fixtures/red.png', 'base64');
+    const image = await parsePng(rawBase64);
+
+    const raw: RawScreenshot = { rawBase64, width: image.width, height: image.height };
+    const decoded: ScreenshotWithOffset = {
+      image,
+      offset: { top: 0, left: 0 },
+      width: image.width,
+      height: image.height,
+    };
+
+    const result = await cropAndCompare(raw, decoded);
+
+    expect(result.isEqual).toBe(true);
+    expect(result.diffPixels).toBe(0);
+  });
+
+  test('detects differences between RawScreenshot and ScreenshotWithOffset', async () => {
+    const redBase64 = fs.readFileSync(__dirname + '/fixtures/red.png', 'base64');
+    const blueBase64 = fs.readFileSync(__dirname + '/fixtures/blue.png', 'base64');
+    const blueImage = await parsePng(blueBase64);
+
+    const raw: RawScreenshot = { rawBase64: redBase64, width: 200, height: 100 };
+    const decoded: ScreenshotWithOffset = {
+      image: blueImage,
+      offset: { top: 0, left: 0 },
+      width: 200,
+      height: 100,
+    };
+
+    const result = await cropAndCompare(raw, decoded);
+
+    expect(result.isEqual).toBe(false);
+    expect(result.diffPixels).toBeGreaterThan(0);
+  });
+
+  test('fast path skips decoding when rawBase64 is identical on ScreenshotWithOffset', async () => {
+    const rawBase64 = fs.readFileSync(__dirname + '/fixtures/blue.png', 'base64');
+    const image = await parsePng(rawBase64);
+
+    const screenshot: ScreenshotWithOffset = {
+      image,
+      offset: { top: 0, left: 0 },
+      width: image.width,
+      height: image.height,
+      rawBase64,
+    };
+
+    const result = await cropAndCompare(screenshot, screenshot);
+
+    expect(result.isEqual).toBe(true);
+    expect(result.diffPixels).toBe(0);
+    expect(result.diffImage).toBeNull();
+  });
+});

From 9db6fb06889470c3b1a7bb86ba52fb2e26fbf98c Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Thu, 25 Jun 2026 17:58:01 +0200
Subject: [PATCH 21/23] Fix tests

---
 test/compare-images.test.ts | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/test/compare-images.test.ts b/test/compare-images.test.ts
index 85b06f7..beebc88 100644
--- a/test/compare-images.test.ts
+++ b/test/compare-images.test.ts
@@ -262,7 +262,7 @@ test('returns isEqual=false when comparing images with 0-size', async () => {
 describe('cropAndCompare with RawScreenshot', () => {
   test('returns isEqual=true when two identical RawScreenshots are compared', async () => {
     const rawBase64 = fs.readFileSync(__dirname + '/fixtures/red.png', 'base64');
-    const screenshot: RawScreenshot = { rawBase64, width: 200, height: 100 };
+    const screenshot: RawScreenshot = { rawBase64, width: 684, height: 116 };
 
     const result = await cropAndCompare(screenshot, screenshot);
 
@@ -277,8 +277,8 @@ describe('cropAndCompare with RawScreenshot', () => {
     const redBase64 = fs.readFileSync(__dirname + '/fixtures/red.png', 'base64');
     const blueBase64 = fs.readFileSync(__dirname + '/fixtures/blue.png', 'base64');
 
-    const first: RawScreenshot = { rawBase64: redBase64, width: 200, height: 100 };
-    const second: RawScreenshot = { rawBase64: blueBase64, width: 200, height: 100 };
+    const first: RawScreenshot = { rawBase64: redBase64, width: 684, height: 116 };
+    const second: RawScreenshot = { rawBase64: blueBase64, width: 684, height: 116 };
 
     const result = await cropAndCompare(first, second);
 
@@ -311,12 +311,12 @@ describe('cropAndCompare with RawScreenshot', () => {
     const blueBase64 = fs.readFileSync(__dirname + '/fixtures/blue.png', 'base64');
     const blueImage = await parsePng(blueBase64);
 
-    const raw: RawScreenshot = { rawBase64: redBase64, width: 200, height: 100 };
+    const raw: RawScreenshot = { rawBase64: redBase64, width: 684, height: 116 };
     const decoded: ScreenshotWithOffset = {
       image: blueImage,
       offset: { top: 0, left: 0 },
-      width: 200,
-      height: 100,
+      width: 684,
+      height: 116,
     };
 
     const result = await cropAndCompare(raw, decoded);

From ae893b4bd50dd64ae0d30ec820a4ea99ac405ec7 Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Fri, 26 Jun 2026 19:05:05 +0200
Subject: [PATCH 22/23] Refactor

---
 test/page-object.test.ts            | 330 ----------------------------
 test/screenshot-page-object.test.ts | 318 ++++++++++++++++++++++++++-
 2 files changed, 317 insertions(+), 331 deletions(-)
 delete mode 100644 test/page-object.test.ts

diff --git a/test/page-object.test.ts b/test/page-object.test.ts
deleted file mode 100644
index 8ef9b27..0000000
--- a/test/page-object.test.ts
+++ /dev/null
@@ -1,330 +0,0 @@
-// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-// SPDX-License-Identifier: Apache-2.0
-import { test, expect, describe, vi } from 'vitest';
-import { ScreenshotPageObject } from '../src/page-objects';
-import useBrowser from '../src/use-browser';
-import './utils/setup-local-driver';
-
-type TestFn = (page: ScreenshotPageObject) => Promise<void>;
-function setupTest(testFn: TestFn, url = './test-page-object.html') {
-  return useBrowser(async browser => {
-    await browser.url(url);
-    await testFn(new ScreenshotPageObject(browser));
-  });
-}
-
-test(
-  'getText',
-  setupTest(async page => {
-    expect(await page.getText('#text-content')).toEqual('Some text');
-  })
-);
-
-test(
-  'getElementsText',
-  setupTest(async page => {
-    expect(await page.getElementsText('#text-content, #scrollable-container')).toEqual([
-      'Some text',
-      'Some scrollable text',
-    ]);
-  })
-);
-
-test(
-  'hoverElement',
-  setupTest(async page => {
-    await page.hoverElement('#hover-button');
-    await page.waitForVisible('#hover-span');
-    expect(await page.getText('#hover-span')).toEqual('Hover success');
-  })
-);
-
-test(
-  'keys',
-  setupTest(async page => {
-    await page.click('#input-1');
-    await page.keys(['Tab']);
-    expect(await page.isFocused('#input-2')).toBe(true);
-  })
-);
-
-test(
-  'isSelected',
-  setupTest(async page => {
-    expect(await page.isSelected('#checkbox')).toBe(true);
-  })
-);
-
-test(
-  'getElementAttribute',
-  setupTest(async page => {
-    expect(await page.getElementAttribute('#checkbox', 'type')).toBe('checkbox');
-  })
-);
-
-test(
-  'getElementProperty',
-  setupTest(async page => {
-    expect(await page.getElementProperty('body', 'tagName')).toBe('BODY');
-  })
-);
-
-test(
-  'getElementsCount',
-  setupTest(async page => {
-    expect(await page.getElementsCount('input')).toBe(3);
-  })
-);
-
-test(
-  'setValue and getValue',
-  setupTest(async page => {
-    expect(await page.getValue('#input-1')).toEqual('');
-    await page.setValue('#input-1', 'test');
-    expect(await page.getValue('#input-1')).toEqual('test');
-  })
-);
-
-test.each([
-  { width: 400, height: 300 },
-  { width: 300, height: 400 },
-])('setWindowSize, width=$width, height=$height', size =>
-  setupTest(async page => {
-    await page.setWindowSize(size);
-    const { width, height } = await page.getViewportSize();
-    expect(width).toBe(size.width);
-
-    // With Chromium --headless=new the window.innerHeight differs from the defined window height.
-    expect(height).toBeGreaterThan(size.height - 100);
-    expect(height).toBeLessThanOrEqual(size.height);
-  })()
-);
-
-test(
-  'spyOnEvents',
-  setupTest(async page => {
-    const spy = await page.spyOnEvents('#button', ['click', 'mouseover']);
-    await page.click('#button');
-    expect(await spy.getEvents()).toEqual(['mouseover', 'click']);
-    await spy.reset();
-    expect(await spy.getEvents()).toEqual([]);
-  })
-);
-
-test(
-  'waitForVisible',
-  setupTest(async page => {
-    await page.waitForVisible('#text-content');
-  })
-);
-
-test(
-  'waitForVisible negated',
-  setupTest(async page => {
-    await page.waitForVisible('#hidden', false);
-  })
-);
-
-test(
-  'waitForExist',
-  setupTest(async page => {
-    await page.waitForExist('#hidden');
-    await page.waitForExist('#not-existing', false);
-  })
-);
-
-describe('waitForAssertion', () => {
-  test(
-    'successful assertion',
-    setupTest(async page => {
-      const assertion = vi.fn(async () => expect(true).toBe(true));
-      await page.waitForAssertion(assertion);
-      expect(assertion).toHaveBeenCalledTimes(1);
-    })
-  );
-
-  test(
-    'retrying once assertion',
-    setupTest(async page => {
-      let counter = 0;
-      const assertion = vi.fn(async () => {
-        counter++;
-        expect(counter).toEqual(2);
-      });
-      await page.waitForAssertion(assertion);
-      expect(assertion).toHaveBeenCalledTimes(2);
-    })
-  );
-
-  test(
-    'reports the original error into the outer scope',
-    setupTest(async page => {
-      const assertion = vi.fn(async () => expect(true).toBe(false));
-      await expect(page.waitForAssertion(assertion)).rejects.toThrowError(/expected true to be false/);
-      expect(assertion).toHaveBeenCalledTimes(6);
-    })
-  );
-});
-
-test(
-  'isExisting',
-  setupTest(async page => {
-    expect(await page.isExisting('#text-content')).toEqual(true);
-    expect(await page.isExisting('#not-existing')).toEqual(false);
-  })
-);
-
-test(
-  'isDisplayed',
-  setupTest(async page => {
-    expect(await page.isDisplayed('#text-content')).toEqual(true);
-    expect(await page.isDisplayed('#text-content-at-page-bottom')).toEqual(true);
-    expect(await page.isDisplayed('#hidden')).toEqual(false);
-    expect(await page.isDisplayed('#not-existing')).toEqual(false);
-  })
-);
-
-test(
-  'isDisplayedInViewport',
-  setupTest(async page => {
-    expect(await page.isDisplayedInViewport('#text-content')).toEqual(true);
-    expect(await page.isDisplayedInViewport('#text-content-at-page-bottom')).toEqual(false);
-    expect(await page.isDisplayedInViewport('#hidden')).toEqual(false);
-    expect(await page.isDisplayedInViewport('#not-existing')).toEqual(false);
-  })
-);
-
-test(
-  'isClickable',
-  setupTest(async page => {
-    await expect(page.isClickable('#hover-button')).resolves.toBe(true);
-    await expect(page.isClickable('#disabled-button')).resolves.toBe(false);
-  })
-);
-
-test(
-  'windowScrollTo/getWindowScroll',
-  setupTest(async page => {
-    await page.windowScrollTo({ top: 40 });
-    expect(await page.getWindowScroll()).toEqual({ top: 40, left: 0 });
-  })
-);
-
-test(
-  'getViewportSize',
-  setupTest(async page => {
-    await expect(page.getViewportSize()).resolves.toEqual({
-      pixelRatio: 1,
-      left: 0,
-      top: 0,
-      width: expect.any(Number),
-      height: expect.any(Number),
-      pageHeight: expect.any(Number),
-      screenHeight: expect.any(Number),
-      screenWidth: expect.any(Number),
-    });
-  })
-);
-
-test(
-  'elementScrollTo/getElementScroll',
-  setupTest(async page => {
-    await page.elementScrollTo('#scrollable-container', { left: 40 });
-    expect(await page.getElementScroll('#scrollable-container')).toEqual({ top: 0, left: 40 });
-  })
-);
-
-test(
-  'elementScrollTo should not scroll when trying to scroll a non-scrollable element',
-  setupTest(async page => {
-    await expect(() => page.elementScrollTo('#text-content', { left: 40 })).rejects.toThrowError(
-      /Element #text-content is not scrollable/
-    );
-  })
-);
-
-test(
-  'elementScrollTo should scroll when one direction is scrollable',
-  setupTest(async page => {
-    await page.elementScrollTo('#vertically-scrollable-container', { top: 40 });
-    expect(await page.getElementScroll('#vertically-scrollable-container')).toEqual({ top: 40, left: 0 });
-  })
-);
-
-test(
-  'elementScrollTo should not scroll in the wrong direction',
-  setupTest(async page => {
-    await expect(() => page.elementScrollTo('#vertically-scrollable-container', { left: 40 })).rejects.toThrowError(
-      / Element #vertically-scrollable-container is not scrollable in left direction/
-    );
-  })
-);
-
-test(
-  'getBoundingBox',
-  setupTest(async page => {
-    const box = await page.getBoundingBox('#text-content');
-    // we can't use absolute numbers in this assertion, because the values are different on Mac and Linux
-    expect(box).toEqual({
-      left: expect.any(Number),
-      right: expect.any(Number),
-      top: expect.any(Number),
-      bottom: expect.any(Number),
-      height: expect.any(Number),
-      width: expect.any(Number),
-    });
-  })
-);
-
-test(
-  'getFocusedElementText',
-  setupTest(async page => {
-    await page.click('#button');
-    expect(await page.getFocusedElementText()).toEqual('Click me');
-  })
-);
-
-test(
-  'click',
-  setupTest(async page => {
-    await page.click('#button');
-    await page.waitForVisible('#click-message');
-  })
-);
-
-test(
-  'click via buttonDown/buttonUp',
-  setupTest(async page => {
-    await page.buttonDownOnElement('#button');
-    await page.buttonUp();
-    await page.waitForVisible('#click-message');
-  })
-);
-
-test(
-  'live announcements',
-  setupTest(async page => {
-    await page.initLiveAnnouncementsObserver();
-    await page.click('#update-live-announcement-button');
-    await page.click('#update-live-announcement-button');
-    await expect(page.getLiveAnnouncements()).resolves.toEqual(['update 1', 'update 2']);
-    await page.clearLiveAnnouncements();
-    await expect(page.getLiveAnnouncements()).resolves.toEqual([]);
-  })
-);
-
-test(
-  'runInsideIframe',
-  setupTest(async page => {
-    await expect(page.isExisting('#inside-iframe')).resolves.toBe(false);
-    await page.runInsideIframe('#test-iframe', true, async () => {
-      await expect(page.isExisting('#inside-iframe')).resolves.toBe(true);
-    });
-    // make sure we exit the iframe properly
-    await expect(page.isExisting('#inside-iframe')).resolves.toBe(false);
-    await page.runInsideIframe('#test-iframe', false, async () => {
-      // should skip switching to iframe here
-      await expect(page.isExisting('#inside-iframe')).resolves.toBe(false);
-    });
-  })
-);
diff --git a/test/screenshot-page-object.test.ts b/test/screenshot-page-object.test.ts
index 6ee4d67..cc34ad0 100644
--- a/test/screenshot-page-object.test.ts
+++ b/test/screenshot-page-object.test.ts
@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import { test, expect, describe } from 'vitest';
+import { test, expect, describe, vi } from 'vitest';
 import { ScreenshotPageObject } from '../src/page-objects';
 import useBrowser from '../src/use-browser';
 import './utils/setup-local-driver';
@@ -13,6 +13,48 @@ function setupTest(testFn: TestFn, url = './test-page-object.html') {
   });
 }
 
+test(
+  'getText',
+  setupTest(async page => {
+    expect(await page.getText('#text-content')).toEqual('Some text');
+  })
+);
+
+test(
+  'getElementsText',
+  setupTest(async page => {
+    expect(await page.getElementsText('#text-content, #scrollable-container')).toEqual([
+      'Some text',
+      'Some scrollable text',
+    ]);
+  })
+);
+
+test(
+  'hoverElement',
+  setupTest(async page => {
+    await page.hoverElement('#hover-button');
+    await page.waitForVisible('#hover-span');
+    expect(await page.getText('#hover-span')).toEqual('Hover success');
+  })
+);
+
+test(
+  'keys',
+  setupTest(async page => {
+    await page.click('#input-1');
+    await page.keys(['Tab']);
+    expect(await page.isFocused('#input-2')).toBe(true);
+  })
+);
+
+test(
+  'isSelected',
+  setupTest(async page => {
+    expect(await page.isSelected('#checkbox')).toBe(true);
+  })
+);
+
 test(
   'focusNextElement',
   setupTest(async page => {
@@ -22,6 +64,211 @@ test(
   })
 );
 
+test(
+  'getElementAttribute',
+  setupTest(async page => {
+    expect(await page.getElementAttribute('#checkbox', 'type')).toBe('checkbox');
+  })
+);
+
+test(
+  'getElementProperty',
+  setupTest(async page => {
+    expect(await page.getElementProperty('body', 'tagName')).toBe('BODY');
+  })
+);
+
+test(
+  'getElementsCount',
+  setupTest(async page => {
+    expect(await page.getElementsCount('input')).toBe(3);
+  })
+);
+
+test(
+  'setValue and getValue',
+  setupTest(async page => {
+    expect(await page.getValue('#input-1')).toEqual('');
+    await page.setValue('#input-1', 'test');
+    expect(await page.getValue('#input-1')).toEqual('test');
+  })
+);
+
+test.each([
+  { width: 400, height: 300 },
+  { width: 300, height: 400 },
+])('setWindowSize, width=$width, height=$height', size =>
+  setupTest(async page => {
+    await page.setWindowSize(size);
+    const { width, height } = await page.getViewportSize();
+    expect(width).toBe(size.width);
+
+    // With Chromium --headless=new the window.innerHeight differs from the defined window height.
+    expect(height).toBeGreaterThan(size.height - 100);
+    expect(height).toBeLessThanOrEqual(size.height);
+  })()
+);
+
+test(
+  'spyOnEvents',
+  setupTest(async page => {
+    const spy = await page.spyOnEvents('#button', ['click', 'mouseover']);
+    await page.click('#button');
+    expect(await spy.getEvents()).toEqual(['mouseover', 'click']);
+    await spy.reset();
+    expect(await spy.getEvents()).toEqual([]);
+  })
+);
+
+test(
+  'waitForVisible',
+  setupTest(async page => {
+    await page.waitForVisible('#text-content');
+  })
+);
+
+test(
+  'waitForVisible negated',
+  setupTest(async page => {
+    await page.waitForVisible('#hidden', false);
+  })
+);
+
+test(
+  'waitForExist',
+  setupTest(async page => {
+    await page.waitForExist('#hidden');
+    await page.waitForExist('#not-existing', false);
+  })
+);
+
+describe('waitForAssertion', () => {
+  test(
+    'successful assertion',
+    setupTest(async page => {
+      const assertion = vi.fn(async () => expect(true).toBe(true));
+      await page.waitForAssertion(assertion);
+      expect(assertion).toHaveBeenCalledTimes(1);
+    })
+  );
+
+  test(
+    'retrying once assertion',
+    setupTest(async page => {
+      let counter = 0;
+      const assertion = vi.fn(async () => {
+        counter++;
+        expect(counter).toEqual(2);
+      });
+      await page.waitForAssertion(assertion);
+      expect(assertion).toHaveBeenCalledTimes(2);
+    })
+  );
+
+  test(
+    'reports the original error into the outer scope',
+    setupTest(async page => {
+      const assertion = vi.fn(async () => expect(true).toBe(false));
+      await expect(page.waitForAssertion(assertion)).rejects.toThrowError(/expected true to be false/);
+      expect(assertion).toHaveBeenCalledTimes(6);
+    })
+  );
+});
+
+test(
+  'isExisting',
+  setupTest(async page => {
+    expect(await page.isExisting('#text-content')).toEqual(true);
+    expect(await page.isExisting('#not-existing')).toEqual(false);
+  })
+);
+
+test(
+  'isDisplayed',
+  setupTest(async page => {
+    expect(await page.isDisplayed('#text-content')).toEqual(true);
+    expect(await page.isDisplayed('#text-content-at-page-bottom')).toEqual(true);
+    expect(await page.isDisplayed('#hidden')).toEqual(false);
+    expect(await page.isDisplayed('#not-existing')).toEqual(false);
+  })
+);
+
+test(
+  'isDisplayedInViewport',
+  setupTest(async page => {
+    expect(await page.isDisplayedInViewport('#text-content')).toEqual(true);
+    expect(await page.isDisplayedInViewport('#text-content-at-page-bottom')).toEqual(false);
+    expect(await page.isDisplayedInViewport('#hidden')).toEqual(false);
+    expect(await page.isDisplayedInViewport('#not-existing')).toEqual(false);
+  })
+);
+
+test(
+  'isClickable',
+  setupTest(async page => {
+    await expect(page.isClickable('#hover-button')).resolves.toBe(true);
+    await expect(page.isClickable('#disabled-button')).resolves.toBe(false);
+  })
+);
+
+test(
+  'windowScrollTo/getWindowScroll',
+  setupTest(async page => {
+    await page.windowScrollTo({ top: 40 });
+    expect(await page.getWindowScroll()).toEqual({ top: 40, left: 0 });
+  })
+);
+
+test(
+  'getViewportSize',
+  setupTest(async page => {
+    await expect(page.getViewportSize()).resolves.toEqual({
+      pixelRatio: 1,
+      left: 0,
+      top: 0,
+      width: expect.any(Number),
+      height: expect.any(Number),
+      pageHeight: expect.any(Number),
+      screenHeight: expect.any(Number),
+      screenWidth: expect.any(Number),
+    });
+  })
+);
+
+test(
+  'elementScrollTo/getElementScroll',
+  setupTest(async page => {
+    await page.elementScrollTo('#scrollable-container', { left: 40 });
+    expect(await page.getElementScroll('#scrollable-container')).toEqual({ top: 0, left: 40 });
+  })
+);
+
+test(
+  'elementScrollTo should not scroll when trying to scroll a non-scrollable element',
+  setupTest(async page => {
+    await expect(() => page.elementScrollTo('#text-content', { left: 40 })).rejects.toThrowError(
+      /Element #text-content is not scrollable/
+    );
+  })
+);
+
+test(
+  'elementScrollTo should scroll when one direction is scrollable',
+  setupTest(async page => {
+    await page.elementScrollTo('#vertically-scrollable-container', { top: 40 });
+    expect(await page.getElementScroll('#vertically-scrollable-container')).toEqual({ top: 40, left: 0 });
+  })
+);
+
+test(
+  'elementScrollTo should not scroll in the wrong direction',
+  setupTest(async page => {
+    await expect(() => page.elementScrollTo('#vertically-scrollable-container', { left: 40 })).rejects.toThrowError(
+      / Element #vertically-scrollable-container is not scrollable in left direction/
+    );
+  })
+);
+
 test(
   'scrollToRight',
   setupTest(async page => {
@@ -47,6 +294,75 @@ test(
   })
 );
 
+test(
+  'getBoundingBox',
+  setupTest(async page => {
+    const box = await page.getBoundingBox('#text-content');
+    // we can't use absolute numbers in this assertion, because the values are different on Mac and Linux
+    expect(box).toEqual({
+      left: expect.any(Number),
+      right: expect.any(Number),
+      top: expect.any(Number),
+      bottom: expect.any(Number),
+      height: expect.any(Number),
+      width: expect.any(Number),
+    });
+  })
+);
+
+test(
+  'getFocusedElementText',
+  setupTest(async page => {
+    await page.click('#button');
+    expect(await page.getFocusedElementText()).toEqual('Click me');
+  })
+);
+
+test(
+  'click',
+  setupTest(async page => {
+    await page.click('#button');
+    await page.waitForVisible('#click-message');
+  })
+);
+
+test(
+  'click via buttonDown/buttonUp',
+  setupTest(async page => {
+    await page.buttonDownOnElement('#button');
+    await page.buttonUp();
+    await page.waitForVisible('#click-message');
+  })
+);
+
+test(
+  'live announcements',
+  setupTest(async page => {
+    await page.initLiveAnnouncementsObserver();
+    await page.click('#update-live-announcement-button');
+    await page.click('#update-live-announcement-button');
+    await expect(page.getLiveAnnouncements()).resolves.toEqual(['update 1', 'update 2']);
+    await page.clearLiveAnnouncements();
+    await expect(page.getLiveAnnouncements()).resolves.toEqual([]);
+  })
+);
+
+test(
+  'runInsideIframe',
+  setupTest(async page => {
+    await expect(page.isExisting('#inside-iframe')).resolves.toBe(false);
+    await page.runInsideIframe('#test-iframe', true, async () => {
+      await expect(page.isExisting('#inside-iframe')).resolves.toBe(true);
+    });
+    // make sure we exit the iframe properly
+    await expect(page.isExisting('#inside-iframe')).resolves.toBe(false);
+    await page.runInsideIframe('#test-iframe', false, async () => {
+      // should skip switching to iframe here
+      await expect(page.isExisting('#inside-iframe')).resolves.toBe(false);
+    });
+  })
+);
+
 describe('capturePermutations', () => {
   test(
     'captures all permutations with correct ids and dimensions',

From 335790b04ef98afe0386ed9673f9509d972382fa Mon Sep 17 00:00:00 2001
From: Joan Perals Tresserra <jotresse@amazon.com>
Date: Fri, 26 Jun 2026 19:06:21 +0200
Subject: [PATCH 23/23] Refactor

---
 test/screenshot-page-object.test.ts | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/test/screenshot-page-object.test.ts b/test/screenshot-page-object.test.ts
index cc34ad0..204a4e2 100644
--- a/test/screenshot-page-object.test.ts
+++ b/test/screenshot-page-object.test.ts
@@ -49,18 +49,18 @@ test(
 );
 
 test(
-  'isSelected',
+  'focusNextElement',
   setupTest(async page => {
-    expect(await page.isSelected('#checkbox')).toBe(true);
+    await page.click('#input-1');
+    await page.focusNextElement();
+    expect(await page.isFocused('#input-2')).toBe(true);
   })
 );
 
 test(
-  'focusNextElement',
+  'isSelected',
   setupTest(async page => {
-    await page.click('#input-1');
-    await page.focusNextElement();
-    expect(await page.isFocused('#input-2')).toBe(true);
+    expect(await page.isSelected('#checkbox')).toBe(true);
   })
 );