File size: 2,696 Bytes
9332ccf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | /**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { Readable } from 'node:stream';
import {
captureHeapSnapshot,
MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES,
} from './memorySnapshot.js';
const { mkdirMock, pipelineMock, getHeapSnapshotMock, createWriteStreamMock } =
vi.hoisted(() => ({
mkdirMock: vi.fn(async () => undefined),
pipelineMock: vi.fn(async () => undefined),
getHeapSnapshotMock: vi.fn(),
createWriteStreamMock: vi.fn(),
}));
vi.mock('node:fs/promises', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:fs/promises')>();
return { ...actual, mkdir: mkdirMock };
});
vi.mock('node:fs', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:fs')>();
return { ...actual, createWriteStream: createWriteStreamMock };
});
vi.mock('node:v8', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:v8')>();
return { ...actual, getHeapSnapshot: getHeapSnapshotMock };
});
vi.mock('node:stream/promises', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:stream/promises')>();
return { ...actual, pipeline: pipelineMock };
});
describe('captureHeapSnapshot', () => {
beforeEach(() => {
mkdirMock.mockClear();
pipelineMock.mockClear();
getHeapSnapshotMock.mockClear().mockReturnValue(Readable.from([]));
createWriteStreamMock
.mockClear()
.mockReturnValue({ write: vi.fn(), end: vi.fn() });
});
afterEach(() => {
vi.clearAllMocks();
});
it('exports the 2 GB auto-capture threshold', () => {
expect(MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES).toBe(2 * 1024 * 1024 * 1024);
});
it('creates the target directory and pipelines the V8 snapshot to disk', async () => {
const target = '/tmp/gemini-test/snapshot.heapsnapshot';
await captureHeapSnapshot(target);
expect(mkdirMock).toHaveBeenCalledWith('/tmp/gemini-test', {
recursive: true,
});
expect(getHeapSnapshotMock).toHaveBeenCalledTimes(1);
expect(createWriteStreamMock).toHaveBeenCalledWith(target);
expect(pipelineMock).toHaveBeenCalledTimes(1);
expect(pipelineMock).toHaveBeenCalledWith(
getHeapSnapshotMock.mock.results[0].value,
createWriteStreamMock.mock.results[0].value,
);
});
it('propagates pipeline failures to the caller', async () => {
pipelineMock.mockRejectedValueOnce(new Error('write failed'));
await expect(
captureHeapSnapshot('/tmp/gemini-test/fail.heapsnapshot'),
).rejects.toThrow('write failed');
});
});
|