| 1 | import type * as Electron from "electron"; |
| 2 | import { assert, beforeEach, expect, it, vi } from "vite-plus/test"; |
| 3 | |
| 4 | import type { ActiveWindow } from "./ActiveWindow.ts"; |
| 5 | import type { RegionSnapShotChild, RegionSnapShotProcess } from "./RegionSnapShot.ts"; |
| 6 | |
| 7 | const { createFromBufferMock, resizeMock } = vi.hoisted(() => ({ |
| 8 | createFromBufferMock: vi.fn(), |
| 9 | resizeMock: vi.fn(), |
| 10 | })); |
| 11 | |
| 12 | vi.mock("electron", () => ({ |
| 13 | nativeImage: { createFromBuffer: createFromBufferMock }, |
| 14 | })); |
| 15 | |
| 16 | import { |
| 17 | captureRegionWindowSnapshot, |
| 18 | makeRegionSnapShotPool, |
| 19 | startRegionSnapShotProcess, |
| 20 | } from "./RegionSnapShot.ts"; |
| 21 | |
| 22 | beforeEach(() => { |
| 23 | vi.clearAllMocks(); |
| 24 | createFromBufferMock.mockReturnValue({ resize: resizeMock }); |
| 25 | }); |
| 26 | |
| 27 | const active = { |
| 28 | title: "Editor", |
| 29 | owner: { name: "Editor", processId: 123 }, |
| 30 | bounds: { x: 10, y: 20, width: 800, height: 600 }, |
| 31 | } as ActiveWindow; |
| 32 | |
| 33 | it("returns a small capture's PNG untouched", async () => { |
| 34 | const png = Buffer.from([1, 2, 3]); |
| 35 | const capture = vi.fn(async () => ({ width: 400, height: 300, png })); |
| 36 | const region = { x: 5, y: 10, width: 400, height: 300 }; |
| 37 | |
| 38 | const result = await captureRegionWindowSnapshot({ capture }, active, region, region); |
| 39 | |
| 40 | assert.deepEqual(capture.mock.calls, [[region]]); |
| 41 | assert.strictEqual(result.png, png); |
| 42 | assert.deepEqual(result.source, { name: "Editor" }); |
| 43 | assert.lengthOf(createFromBufferMock.mock.calls, 0); |
| 44 | }); |
| 45 | |
| 46 | it.each([ |
| 47 | { width: 600, height: 400, expected: { width: 240, height: 160 } }, |
| 48 | { width: 400, height: 600, expected: { width: 107, height: 160 } }, |
| 49 | { width: 600, height: 1, expected: { width: 256, height: 1 } }, |
| 50 | ])( |
| 51 | "bounds $width x $height physical pixels before encoding", |
| 52 | async ({ width, height, expected }) => { |
| 53 | const png = Buffer.from([1, 2, 3]); |
| 54 | const resizedPng = Buffer.from([4, 5, 6]); |
| 55 | const capture = vi.fn(async () => ({ width, height, png })); |
| 56 | resizeMock.mockReturnValue({ toPNG: () => resizedPng }); |
| 57 | |
| 58 | const result = await captureRegionWindowSnapshot({ capture }, active, active.bounds, { |
| 59 | width: 256, |
| 60 | height: 160, |
| 61 | }); |
| 62 | |
| 63 | assert.deepEqual(createFromBufferMock.mock.calls, [[png]]); |
| 64 | assert.deepEqual(resizeMock.mock.calls, [[{ ...expected, quality: "best" }]]); |
| 65 | assert.strictEqual(result.png, resizedPng); |
| 66 | }, |
| 67 | ); |
| 68 | |
| 69 | function fakeChild() { |
| 70 | const messageListeners: Array<(message: unknown) => void> = []; |
| 71 | const exitListeners: Array<() => void> = []; |
| 72 | const send = vi.fn(); |
| 73 | const kill = vi.fn(); |
| 74 | const child: RegionSnapShotChild = { |
| 75 | send, |
| 76 | onMessage: (listener) => messageListeners.push(listener), |
| 77 | onExit: (listener) => exitListeners.push(listener), |
| 78 | kill, |
| 79 | }; |
| 80 | return { |
| 81 | child, |
| 82 | send, |
| 83 | kill, |
| 84 | emit: (message: unknown) => { |
| 85 | for (const listener of messageListeners) listener(message); |
| 86 | }, |
| 87 | exit: () => { |
| 88 | for (const listener of exitListeners) listener(); |
| 89 | }, |
| 90 | }; |
| 91 | } |
| 92 | |
| 93 | const region = { x: 0, y: 0, width: 2, height: 1 }; |
| 94 | |
| 95 | it("resolves the decoded capture once the child is ready and replies", async () => { |
| 96 | const fake = fakeChild(); |
| 97 | const fork = vi.fn(() => fake.child); |
| 98 | const png = Buffer.from([9, 8, 7]); |
| 99 | const child = startRegionSnapShotProcess("worker.cjs", fork); |
| 100 | assert.deepEqual(fork.mock.calls, [["worker.cjs"]]); |
| 101 | assert.isTrue(child.alive); |
| 102 | |
| 103 | const capture = child.capture(region); |
| 104 | await Promise.resolve(); |
| 105 | assert.lengthOf(fake.send.mock.calls, 0); |
| 106 | fake.emit("ready"); |
| 107 | await vi.waitFor(() => assert.deepEqual(fake.send.mock.calls, [[{ region }]])); |
| 108 | fake.emit({ type: "result", width: 2, height: 1, png: png.toString("base64") }); |
| 109 | |
| 110 | assert.deepEqual(await capture, { width: 2, height: 1, png }); |
| 111 | assert.isFalse(child.alive); |
| 112 | assert.lengthOf(fake.kill.mock.calls, 1); |
| 113 | }); |
| 114 | |
| 115 | it("rejects with the child's error message", async () => { |
| 116 | const fake = fakeChild(); |
| 117 | const child = startRegionSnapShotProcess("worker.cjs", () => fake.child); |
| 118 | const capture = child.capture(region); |
| 119 | fake.emit("ready"); |
| 120 | fake.emit({ type: "error", message: "Screen capture failed" }); |
| 121 | |
| 122 | await expect(capture).rejects.toThrow("Screen capture failed"); |
| 123 | assert.isFalse(child.alive); |
| 124 | }); |
| 125 | |
| 126 | it("rejects when the child exits before replying", async () => { |
| 127 | const fake = fakeChild(); |
| 128 | const child = startRegionSnapShotProcess("worker.cjs", () => fake.child); |
| 129 | const capture = child.capture(region); |
| 130 | fake.emit("ready"); |
| 131 | fake.exit(); |
| 132 | |
| 133 | await expect(capture).rejects.toThrow(/helper exited/); |
| 134 | }); |
| 135 | |
| 136 | it("serves a single capture per child", async () => { |
| 137 | const fake = fakeChild(); |
| 138 | const child = startRegionSnapShotProcess("worker.cjs", () => fake.child); |
| 139 | const first = child.capture(region); |
| 140 | fake.emit("ready"); |
| 141 | await vi.waitFor(() => assert.lengthOf(fake.send.mock.calls, 1)); |
| 142 | fake.emit({ type: "result", width: 1, height: 1, png: Buffer.from([1]).toString("base64") }); |
| 143 | await first; |
| 144 | |
| 145 | await expect(child.capture(region)).rejects.toThrow(/already used/); |
| 146 | assert.lengthOf(fake.send.mock.calls, 1); |
| 147 | }); |
| 148 | |
| 149 | it("cancels a pending capture on close and kills the child", async () => { |
| 150 | const fake = fakeChild(); |
| 151 | const child = startRegionSnapShotProcess("worker.cjs", () => fake.child); |
| 152 | const capture = child.capture(region); |
| 153 | child.close(); |
| 154 | |
| 155 | await expect(capture).rejects.toThrow(/cancelled/); |
| 156 | assert.lengthOf(fake.kill.mock.calls, 1); |
| 157 | assert.isFalse(child.alive); |
| 158 | fake.emit("ready"); |
| 159 | assert.lengthOf(fake.send.mock.calls, 0); |
| 160 | }); |
| 161 | |
| 162 | it("reports an unavailable helper when the child cannot be forked", async () => { |
| 163 | const child = startRegionSnapShotProcess("worker.cjs", () => { |
| 164 | throw new Error("fork failed"); |
| 165 | }); |
| 166 | |
| 167 | assert.isFalse(child.alive); |
| 168 | await expect(child.capture(region)).rejects.toThrow(/unavailable/); |
| 169 | }); |
| 170 | |
| 171 | it("times out a stalled capture", async () => { |
| 172 | vi.useFakeTimers(); |
| 173 | try { |
| 174 | const fake = fakeChild(); |
| 175 | const child = startRegionSnapShotProcess("worker.cjs", () => fake.child); |
| 176 | const failure = expect(child.capture(region)).rejects.toThrow(/timed out/); |
| 177 | fake.emit("ready"); |
| 178 | await vi.advanceTimersByTimeAsync(15_000); |
| 179 | |
| 180 | await failure; |
| 181 | assert.lengthOf(fake.kill.mock.calls, 1); |
| 182 | } finally { |
| 183 | vi.useRealTimers(); |
| 184 | } |
| 185 | }); |
| 186 | |
| 187 | function fakeProcess() { |
| 188 | const result = Promise.withResolvers<{ width: number; height: number; png: Buffer }>(); |
| 189 | void result.promise.catch(() => undefined); |
| 190 | let alive = true; |
| 191 | const close = vi.fn(() => { |
| 192 | alive = false; |
| 193 | result.reject(new Error("Windows window capture cancelled.")); |
| 194 | }); |
| 195 | const capture = vi.fn((_region: Electron.Rectangle) => result.promise); |
| 196 | const process: RegionSnapShotProcess = { |
| 197 | get alive() { |
| 198 | return alive; |
| 199 | }, |
| 200 | capture, |
| 201 | close, |
| 202 | }; |
| 203 | return { |
| 204 | process, |
| 205 | capture, |
| 206 | close, |
| 207 | resolve: (value: { width: number; height: number; png: Buffer }) => { |
| 208 | alive = false; |
| 209 | result.resolve(value); |
| 210 | }, |
| 211 | }; |
| 212 | } |
| 213 | |
| 214 | function fakePool() { |
| 215 | const processes: Array<ReturnType<typeof fakeProcess>> = []; |
| 216 | const start = vi.fn(() => { |
| 217 | const process = fakeProcess(); |
| 218 | processes.push(process); |
| 219 | return process.process; |
| 220 | }); |
| 221 | return { pool: makeRegionSnapShotPool("worker.cjs", start), processes, start }; |
| 222 | } |
| 223 | |
| 224 | it("warms one standby child and keeps it while alive", () => { |
| 225 | const { pool, start } = fakePool(); |
| 226 | |
| 227 | pool.warm(); |
| 228 | pool.warm(); |
| 229 | |
| 230 | assert.deepEqual(start.mock.calls, [["worker.cjs"]]); |
| 231 | }); |
| 232 | |
| 233 | it("captures with the standby child and warms a replacement afterwards", async () => { |
| 234 | const { pool, processes, start } = fakePool(); |
| 235 | pool.warm(); |
| 236 | const shot = { width: 1, height: 1, png: Buffer.from([1]) }; |
| 237 | |
| 238 | const capture = pool.capture(region); |
| 239 | assert.deepEqual(processes[0]!.capture.mock.calls, [[region]]); |
| 240 | assert.lengthOf(start.mock.calls, 1); |
| 241 | processes[0]!.resolve(shot); |
| 242 | |
| 243 | assert.deepEqual(await capture, shot); |
| 244 | assert.lengthOf(start.mock.calls, 2); |
| 245 | assert.lengthOf(processes[1]!.capture.mock.calls, 0); |
| 246 | }); |
| 247 | |
| 248 | it("starts a child on demand when no standby is warm", async () => { |
| 249 | const { pool, processes, start } = fakePool(); |
| 250 | const shot = { width: 1, height: 1, png: Buffer.from([1]) }; |
| 251 | |
| 252 | const capture = pool.capture(region); |
| 253 | assert.lengthOf(start.mock.calls, 1); |
| 254 | processes[0]!.resolve(shot); |
| 255 | |
| 256 | assert.deepEqual(await capture, shot); |
| 257 | }); |
| 258 | |
| 259 | it("rejects a second capture while one is in flight", async () => { |
| 260 | const { pool, processes } = fakePool(); |
| 261 | const first = pool.capture(region); |
| 262 | |
| 263 | await expect(pool.capture(region)).rejects.toThrow(/still in progress/); |
| 264 | assert.lengthOf(processes, 1); |
| 265 | |
| 266 | processes[0]!.resolve({ width: 1, height: 1, png: Buffer.from([1]) }); |
| 267 | await first; |
| 268 | }); |
| 269 | |
| 270 | it("cools the standby child", () => { |
| 271 | const { pool, processes, start } = fakePool(); |
| 272 | pool.warm(); |
| 273 | |
| 274 | pool.cool(); |
| 275 | |
| 276 | assert.lengthOf(processes[0]!.close.mock.calls, 1); |
| 277 | pool.warm(); |
| 278 | assert.lengthOf(start.mock.calls, 2); |
| 279 | }); |
| 280 | |
| 281 | it("closes the standby and the in-flight child", async () => { |
| 282 | const { pool, processes } = fakePool(); |
| 283 | pool.warm(); |
| 284 | const capture = pool.capture(region); |
| 285 | pool.warm(); |
| 286 | assert.lengthOf(processes, 2); |
| 287 | |
| 288 | pool.close(); |
| 289 | |
| 290 | await expect(capture).rejects.toThrow(/cancelled/); |
| 291 | assert.lengthOf(processes[0]!.close.mock.calls, 1); |
| 292 | assert.lengthOf(processes[1]!.close.mock.calls, 1); |
| 293 | await expect(pool.capture(region)).rejects.toThrow(/unavailable/); |
| 294 | assert.lengthOf(processes, 2); |
| 295 | }); |