fix(screenshot): white sampling (#636)

* fix(screenshot): white sampling

* chore(screenshot): handle max memory limit

* refactor: better jimp helper

Author: Kikobeats

packages/screenshot/src/is-white-screenshot.js

@@ -3,20 +3,30 @@
 const { Jimp } = require('jimp')
 
 module.exports = async uint8array => {
-  const image = await Jimp.fromBuffer(Buffer.from(uint8array))
-  const firstPixel = image.getPixelColor(0, 0)
-  const height = image.bitmap.height
-  const width = image.bitmap.width
+  try {
+    const image = await Jimp.fromBuffer(Buffer.from(uint8array))
+    const firstPixel = image.getPixelColor(0, 0)
+    const height = image.bitmap.height
+    const width = image.bitmap.width
 
-  const samplePercentage = 0.25 // Sample 25% of the image
-  const sampleSize = Math.floor(width * height * samplePercentage) // Calculate sample size based on percentage
-  const stepSize = Math.max(1, Math.floor((width * height) / sampleSize)) // Calculate step size based on sample size
+    // For 2D grid sampling, calculate stepSize to achieve approximately the target sample percentage.
+    // When sampling every 'stepSize' pixels in both dimensions, actual samples = (height/stepSize) * (width/stepSize).
+    // To achieve samplePercentage, we need: (h*w)/(stepSize²) ≈ samplePercentage*(h*w)
+    // Therefore: stepSize ≈ sqrt(1 / samplePercentage)
+    const samplePercentage = 0.25 // Sample ~25% of the image
+    const stepSize = Math.max(1, Math.ceil(Math.sqrt(1 / samplePercentage)))
 
-  for (let i = 0; i < height; i += stepSize) {
-    for (let j = 0; j < width; j += stepSize) {
-      if (firstPixel !== image.getPixelColor(j, i)) return false
+    for (let i = 0; i < height; i += stepSize) {
+      for (let j = 0; j < width; j += stepSize) {
+        if (firstPixel !== image.getPixelColor(j, i)) return false
+      }
     }
-  }
 
-  return true
+    return true
+  } catch (error) {
+    if (error.message.includes('maxMemoryUsageInMB')) {
+      return false
+    }
+    throw error
+  }
 }

packages/screenshot/test/is-white-screenshot.js

@@ -3,8 +3,36 @@
 const test = require('ava')
 const { readFile } = require('fs/promises')
 
+const { Jimp } = require('jimp')
+
 const isWhite = require('../src/is-white-screenshot')
 
+const createJimpSpy = () => {
+  const originalFromBuffer = Jimp.fromBuffer
+  const spy = { callCount: 0 }
+
+  const wrappedFromBuffer = async function (buffer, options) {
+    const image = await originalFromBuffer.call(this, buffer, options)
+    const originalGetPixelColor = image.getPixelColor.bind(image)
+
+    image.getPixelColor = function (x, y) {
+      spy.callCount++
+      return originalGetPixelColor(x, y)
+    }
+
+    return image
+  }
+
+  Jimp.fromBuffer = wrappedFromBuffer
+
+  return {
+    spy,
+    restore: () => {
+      Jimp.fromBuffer = originalFromBuffer
+    }
+  }
+}
+
 test('true', async t => {
   t.true(await isWhite(await readFile('./test/fixtures/white-5k.jpg')))
   t.true(await isWhite(await readFile('./test/fixtures/white-5k.png')))
@@ -14,3 +42,38 @@ test('false', async t => {
   t.false(await isWhite(await readFile('./test/fixtures/no-white-5k.jpg')))
   t.false(await isWhite(await readFile('./test/fixtures/no-white-5k.png')))
 })
+
+test('sampling algorithm correctly samples ~25% of pixels', async t => {
+  const imageBuffer = await readFile('./test/fixtures/white-5k.png')
+  const tempImage = await Jimp.fromBuffer(imageBuffer)
+  const totalPixels = tempImage.bitmap.width * tempImage.bitmap.height
+  const { spy, restore } = createJimpSpy()
+
+  await isWhite(imageBuffer)
+  restore()
+
+  const percentageChecked = (spy.callCount / totalPixels) * 100
+
+  t.true(
+    percentageChecked >= 20 && percentageChecked <= 30,
+    `Expected to check ~25% of pixels, but checked ${percentageChecked.toFixed(2)}% (${
+      spy.callCount
+    }/${totalPixels})`
+  )
+})
+
+test('handles memory errors gracefully on very large images', async t => {
+  const { createCanvas } = require('canvas')
+
+  const width = 10000
+  const height = 10000
+  const canvas = createCanvas(width, height)
+  const ctx = canvas.getContext('2d')
+
+  ctx.fillStyle = '#ffffff'
+  ctx.fillRect(0, 0, width, height)
+
+  const largeImageBuffer = canvas.toBuffer('image/jpeg', { quality: 0.9 })
+  const result = await isWhite(largeImageBuffer)
+  t.is(typeof result, 'boolean', 'Should return a boolean, not throw an error')
+})