From f1faca416ead45b381deb86ffc5a2c2b2c74fbf4 Mon Sep 17 00:00:00 2001
From: Fedor Indutny <79877362+indutny-signal@users.noreply.github.com>
Date: Thu, 28 Aug 2025 09:44:44 -0700
Subject: [PATCH] Add `.screenshot()` method to mock tests

---
 ts/test-mock/bootstrap.ts | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/ts/test-mock/bootstrap.ts b/ts/test-mock/bootstrap.ts
index 1e5182d37f..61fdcaef23 100644
--- a/ts/test-mock/bootstrap.ts
+++ b/ts/test-mock/bootstrap.ts
@@ -221,6 +221,7 @@ export class Bootstrap {
   #storagePath?: string;
   #timestamp: number = Date.now() - WEEK;
   #lastApp?: App;
+  #screenshotId = 0;
   readonly #randomId = crypto.randomBytes(8).toString('hex');
 
   constructor(options: BootstrapOptions = {}) {
@@ -545,6 +546,28 @@ export class Bootstrap {
     }
   }
 
+  public async screenshot(
+    app: App | undefined = this.#lastApp,
+    testName?: string
+  ): Promise<void> {
+    if (!app) {
+      return;
+    }
+
+    const outDir = await this.#getArtifactsDir(testName);
+    if (outDir == null) {
+      return;
+    }
+
+    const window = await app.getWindow();
+    const screenshot = await window.screenshot();
+
+    const id = this.#screenshotId;
+    this.#screenshotId += 1;
+
+    await fs.writeFile(path.join(outDir, `screenshot-${id}.png`), screenshot);
+  }
+
   public async saveLogs(
     app: App | undefined = this.#lastApp,
     testName?: string
@@ -566,11 +589,7 @@ export class Bootstrap {
         ?.context()
         .tracing.stop({ path: path.join(outDir, 'trace.zip') });
     }
-    if (app) {
-      const window = await app.getWindow();
-      const screenshot = await window.screenshot();
-      await fs.writeFile(path.join(outDir, 'screenshot.png'), screenshot);
-    }
+    await this.screenshot(app, testName);
   }
 
   public async createScreenshotComparator(