/** * @license * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { RestoreCommand, ListCheckpointsCommand } from './restore.js'; import * as fs from 'node:fs/promises'; import { getCheckpointInfoList, getToolCallDataSchema, isNodeError, performRestore, } from '@google/gemini-cli-core'; import type { CommandContext } from './types.js'; import type { Mock } from 'vitest'; vi.mock('node:fs/promises'); vi.mock('@google/gemini-cli-core', async (importOriginal) => { const actual = await importOriginal(); return { ...actual, getCheckpointInfoList: vi.fn(), getToolCallDataSchema: vi.fn(), isNodeError: vi.fn(), performRestore: vi.fn(), }; }); describe('RestoreCommand', () => { let context: CommandContext; let restoreCommand: RestoreCommand; beforeEach(() => { vi.resetAllMocks(); restoreCommand = new RestoreCommand(); context = { agentContext: { config: { getCheckpointingEnabled: vi.fn().mockReturnValue(true), storage: { getProjectTempCheckpointsDir: vi .fn() .mockReturnValue('/tmp/checkpoints'), }, }, }, git: {}, sendMessage: vi.fn(), } as unknown as CommandContext; }); it('delegates to list behavior when invoked without args', async () => { const listExecuteSpy = vi .spyOn(ListCheckpointsCommand.prototype, 'execute') .mockResolvedValue({ name: 'restore list', data: 'list data', }); const response = await restoreCommand.execute(context, []); expect(listExecuteSpy).toHaveBeenCalledWith(context); expect(response).toEqual({ name: 'restore list', data: 'list data', }); }); it('returns checkpointing-disabled message when disabled', async () => { ( context.agentContext.config.getCheckpointingEnabled as Mock ).mockReturnValue(false); const response = await restoreCommand.execute(context, ['checkpoint1']); expect(response.data).toContain('Checkpointing is not enabled'); }); it('returns file-not-found message for missing checkpoint', async () => { const error = new Error('ENOENT'); (error as Error & { code: string }).code = 'ENOENT'; vi.mocked(fs.readFile).mockRejectedValue(error); vi.mocked(isNodeError).mockReturnValue(true); const response = await restoreCommand.execute(context, ['missing']); expect(response.data).toBe('File not found: missing.json'); }); it('handles checkpoint filename already ending in .json', async () => { const error = new Error('ENOENT'); (error as Error & { code: string }).code = 'ENOENT'; vi.mocked(fs.readFile).mockRejectedValue(error); vi.mocked(isNodeError).mockReturnValue(true); const response = await restoreCommand.execute(context, ['existing.json']); expect(response.data).toBe('File not found: existing.json'); expect(fs.readFile).toHaveBeenCalledWith( expect.stringContaining('existing.json'), 'utf-8', ); }); it('returns invalid/corrupt checkpoint message when schema parse fails', async () => { vi.mocked(fs.readFile).mockResolvedValue('{"invalid": "data"}'); vi.mocked(getToolCallDataSchema).mockReturnValue({ safeParse: vi.fn().mockReturnValue({ success: false }), } as unknown as ReturnType); const response = await restoreCommand.execute(context, ['invalid']); expect(response.data).toBe('Checkpoint file is invalid or corrupted.'); }); it('formats streamed restore results correctly', async () => { vi.mocked(fs.readFile).mockResolvedValue('{"valid": "data"}'); vi.mocked(getToolCallDataSchema).mockReturnValue({ safeParse: vi .fn() .mockReturnValue({ success: true, data: { some: 'data' } }), } as unknown as ReturnType); async function* mockRestoreGenerator() { yield { type: 'message', messageType: 'info', content: 'Restoring...' }; yield { type: 'load_history', clientHistory: [{}, {}] }; yield { type: 'other', some: 'other' }; } vi.mocked(performRestore).mockReturnValue( mockRestoreGenerator() as unknown as ReturnType, ); const response = await restoreCommand.execute(context, ['valid']); expect(response.data).toContain('[INFO] Restoring...'); expect(response.data).toContain('Loaded history with 2 messages.'); expect(response.data).toContain( 'Restored: {"type":"other","some":"other"}', ); }); it('returns generic unexpected error message for non-ENOENT failures', async () => { vi.mocked(fs.readFile).mockRejectedValue(new Error('Random error')); vi.mocked(isNodeError).mockReturnValue(false); const response = await restoreCommand.execute(context, ['error']); expect(response.data).toContain( 'An unexpected error occurred during restore: Error: Random error', ); }); }); describe('ListCheckpointsCommand', () => { let context: CommandContext; let listCommand: ListCheckpointsCommand; let mockReaddir: Mock<(path: string) => Promise>; beforeEach(() => { vi.resetAllMocks(); listCommand = new ListCheckpointsCommand(); mockReaddir = vi.mocked(fs.readdir) as unknown as Mock< (path: string) => Promise >; context = { agentContext: { config: { getCheckpointingEnabled: vi.fn().mockReturnValue(true), storage: { getProjectTempCheckpointsDir: vi .fn() .mockReturnValue('/tmp/checkpoints'), }, }, }, } as unknown as CommandContext; }); it('returns checkpointing-disabled message when disabled', async () => { ( context.agentContext.config.getCheckpointingEnabled as Mock ).mockReturnValue(false); const response = await listCommand.execute(context); expect(response.data).toContain('Checkpointing is not enabled'); }); it('returns "No checkpoints found." when no .json checkpoints exist', async () => { mockReaddir.mockResolvedValue(['not-a-checkpoint.txt']); const response = await listCommand.execute(context); expect(response.data).toBe('No checkpoints found.'); }); it('ignores error when mkdir fails', async () => { vi.mocked(fs.mkdir).mockRejectedValue(new Error('mkdir fail')); mockReaddir.mockResolvedValue([]); const response = await listCommand.execute(context); expect(response.data).toBe('No checkpoints found.'); expect(fs.mkdir).toHaveBeenCalled(); }); it('formats checkpoint summary output from checkpoint metadata', async () => { mockReaddir.mockResolvedValue(['cp1.json', 'cp2.json']); vi.mocked(getCheckpointInfoList).mockReturnValue([ { messageId: 'id1', checkpoint: 'cp1' }, { messageId: 'id2', checkpoint: 'cp2' }, ]); const response = await listCommand.execute(context); expect(response.data).toContain('Available Checkpoints:'); // Note: The current implementation of ListCheckpointsCommand incorrectly accesses // fileName, toolName, etc. which don't exist on CheckpointInfo, resulting in 'Unknown'. expect(response.data).toContain('- **Unknown**: Unknown (Status: Unknown)'); }); it('handles empty checkpoint info list', async () => { mockReaddir.mockResolvedValue(['some.json']); vi.mocked(getCheckpointInfoList).mockReturnValue([]); const response = await listCommand.execute(context); expect(response.data).toBe('Available Checkpoints:\n'); }); it('returns generic unexpected error message on failures', async () => { mockReaddir.mockRejectedValue(new Error('Readdir fail')); const response = await listCommand.execute(context); expect(response.data).toBe( 'An unexpected error occurred while listing checkpoints.', ); }); });