File size: 5,836 Bytes
7a1ad33 | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | /**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { performRestore } from './restore.js';
import { type ToolCallData } from '../utils/checkpointUtils.js';
import type { GitService } from '../services/gitService.js';
describe('performRestore', () => {
let mockGitService: GitService;
beforeEach(() => {
mockGitService = {
initialize: vi.fn(),
verifyGitAvailability: vi.fn(),
setupShadowGitRepository: vi.fn(),
getCurrentCommitHash: vi.fn(),
createFileSnapshot: vi.fn(),
restoreProjectFromSnapshot: vi.fn(),
storage: {},
getHistoryDir: vi.fn().mockReturnValue('mock-history-dir'),
shadowGitRepository: {},
} as unknown as GitService;
});
it('should yield load_history if history and clientHistory are present', async () => {
const toolCallData: ToolCallData = {
toolCall: { name: 'test', args: {} },
history: [{ some: 'history' }],
clientHistory: [{ role: 'user', parts: [{ text: 'hello' }] }],
};
const generator = performRestore(toolCallData, undefined);
const result = await generator.next();
expect(result.value).toEqual({
type: 'load_history',
history: toolCallData.history,
clientHistory: toolCallData.clientHistory,
});
expect(result.done).toBe(false);
const nextResult = await generator.next();
expect(nextResult.done).toBe(true);
});
it('should call restoreProjectFromSnapshot and yield a message if commitHash and gitService are present', async () => {
const toolCallData: ToolCallData = {
toolCall: { name: 'test', args: {} },
commitHash: 'test-commit-hash',
};
const spy = vi
.spyOn(mockGitService, 'restoreProjectFromSnapshot')
.mockResolvedValue(undefined);
const generator = performRestore(toolCallData, mockGitService);
const result = await generator.next();
expect(spy).toHaveBeenCalledWith('test-commit-hash');
expect(result.value).toEqual({
type: 'message',
messageType: 'info',
content: 'Restored project to the state before the tool call.',
});
expect(result.done).toBe(false);
const nextResult = await generator.next();
expect(nextResult.done).toBe(true);
});
it('should yield an error message if restoreProjectFromSnapshot throws "unable to read tree" error', async () => {
const toolCallData: ToolCallData = {
toolCall: { name: 'test', args: {} },
commitHash: 'invalid-commit-hash',
};
const spy = vi
.spyOn(mockGitService, 'restoreProjectFromSnapshot')
.mockRejectedValue(
new Error('fatal: unable to read tree invalid-commit-hash'),
);
const generator = performRestore(toolCallData, mockGitService);
const result = await generator.next();
expect(spy).toHaveBeenCalledWith('invalid-commit-hash');
expect(result.value).toEqual({
type: 'message',
messageType: 'error',
content:
"The commit hash 'invalid-commit-hash' associated with this checkpoint could not be found in your Git repository. This can happen if the repository has been re-cloned, reset, or if old commits have been garbage collected. This checkpoint cannot be restored.",
});
expect(result.done).toBe(false);
const nextResult = await generator.next();
expect(nextResult.done).toBe(true);
});
it('should re-throw other errors from restoreProjectFromSnapshot', async () => {
const toolCallData: ToolCallData = {
toolCall: { name: 'test', args: {} },
commitHash: 'some-commit-hash',
};
const testError = new Error('something went wrong');
vi.spyOn(mockGitService, 'restoreProjectFromSnapshot').mockRejectedValue(
testError,
);
const generator = performRestore(toolCallData, mockGitService);
await expect(generator.next()).rejects.toThrow(testError);
});
it('should yield load_history then a message if both are present', async () => {
const toolCallData: ToolCallData = {
toolCall: { name: 'test', args: {} },
history: [{ some: 'history' }],
clientHistory: [{ role: 'user', parts: [{ text: 'hello' }] }],
commitHash: 'test-commit-hash',
};
const spy = vi
.spyOn(mockGitService, 'restoreProjectFromSnapshot')
.mockResolvedValue(undefined);
const generator = performRestore(toolCallData, mockGitService);
const historyResult = await generator.next();
expect(historyResult.value).toEqual({
type: 'load_history',
history: toolCallData.history,
clientHistory: toolCallData.clientHistory,
});
expect(historyResult.done).toBe(false);
const messageResult = await generator.next();
expect(spy).toHaveBeenCalledWith('test-commit-hash');
expect(messageResult.value).toEqual({
type: 'message',
messageType: 'info',
content: 'Restored project to the state before the tool call.',
});
expect(messageResult.done).toBe(false);
const nextResult = await generator.next();
expect(nextResult.done).toBe(true);
});
it('should yield error message if commitHash is present but gitService is undefined', async () => {
const toolCallData: ToolCallData = {
toolCall: { name: 'test', args: {} },
commitHash: 'test-commit-hash',
};
const generator = performRestore(toolCallData, undefined);
const result = await generator.next();
expect(result.value).toEqual({
type: 'message',
messageType: 'error',
content:
'Git service is not available, cannot restore checkpoint. Please ensure you are in a git repository.',
});
expect(result.done).toBe(false);
const nextResult = await generator.next();
expect(nextResult.done).toBe(true);
});
});
|