feat(scenario): "take-screenshot" command

Can either save the screenshot to a file, or compare it with an existing PNG file on disk

Author: urish

package-lock.json

@@ -41,10 +41,12 @@
     "arg": "^5.0.2",
     "chalk": "^5.3.0",
     "chalk-template": "^1.1.0",
+    "pngjs": "^7.0.0",
     "ws": "^8.13.0",
     "yaml": "^2.3.1"
   },
   "devDependencies": {
+    "@types/pngjs": "^6.0.5",
     "@types/ws": "^8.5.4",
     "@typescript-eslint/eslint-plugin": "^5.59.1",
     "@yao-pkg/pkg": "^5.11.5",

package.json

@@ -20,6 +20,7 @@ import { SetControlCommand } from './scenario/SetControlCommand.js';
 import { WaitSerialCommand } from './scenario/WaitSerialCommand.js';
 import { WriteSerialCommand } from './scenario/WriteSerialCommand.js';
 import { uploadFirmware } from './uploadFirmware.js';
+import { TakeScreenshotCommand } from './scenario/TakeScreenshotCommand.js';
 
 const millis = 1_000_000;
 
@@ -215,6 +216,7 @@ async function main() {
       'set-control': new SetControlCommand(),
       'wait-serial': new WaitSerialCommand(expectEngine),
       'write-serial': new WriteSerialCommand(),
+      'take-screenshot': new TakeScreenshotCommand(path.dirname(resolvedScenarioFile)),
     });
     scenario.validate();
   }

src/main.ts

@@ -0,0 +1,68 @@
+import { readFile, writeFile } from 'fs/promises';
+import path from 'path';
+import { PNG } from 'pngjs';
+import { type APIClient } from '../APIClient.js';
+import { type TestScenario } from '../TestScenario.js';
+
+export interface ITakeScreenshotParams {
+  'part-id': string;
+  'save-to'?: string;
+  'compare-with'?: string;
+}
+
+export class TakeScreenshotCommand {
+  constructor(readonly scenarioDir: string) {}
+
+  validate(value: ITakeScreenshotParams) {
+    if (!value['part-id']) {
+      throw new Error('take-screenshot: `part-id` is required');
+    }
+    if (!value['save-to'] && !value['compare-with']) {
+      throw new Error('take-screenshot: `save-to` or `compare-with` is required');
+    }
+  }
+
+  async run(scenario: TestScenario, client: APIClient, params: ITakeScreenshotParams) {
+    const framebuffer = await client.framebufferRead(params['part-id']);
+    const png = Buffer.from(framebuffer.png, 'base64');
+    const saveTo = params['save-to'];
+    const compareWith = params['compare-with'];
+    if (saveTo) {
+      await writeFile(path.join(this.scenarioDir, saveTo), png);
+    }
+    if (compareWith) {
+      const compareWithPath = path.join(this.scenarioDir, compareWith);
+      const originalPng = PNG.sync.read(png);
+
+      let compareWithData;
+      try {
+        compareWithData = await readFile(compareWithPath);
+      } catch (error) {
+        scenario.fail(`Failed to read comparison file for screenshot: ${compareWith}`);
+        return;
+      }
+
+      const compareWithPng = PNG.sync.read(compareWithData);
+
+      if (
+        compareWithPng.width !== originalPng.width ||
+        compareWithPng.height !== originalPng.height
+      ) {
+        scenario.fail(
+          `Comparison file for screenshot (${compareWith}) has a different size (${compareWithPng.width}x${compareWithPng.height}) than the part's screen (${originalPng.width}x${originalPng.height})`,
+        );
+        return;
+      }
+
+      for (let i = 0; i < originalPng.data.length; i++) {
+        const original = originalPng.data[i];
+        const compare = compareWithPng.data[i];
+        if (original !== compare) {
+          const x = (i >> 2) % originalPng.width;
+          const y = Math.floor((i >> 2) / originalPng.width);
+          scenario.fail(`Screenshot mismatch for ${compareWith} at x,y=(${x},${y})`);
+        }
+      }
+    }
+  }
+}