diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts index 47f1faf55..489706978 100644 --- a/packages/backend/src/managers/modelsManager.spec.ts +++ b/packages/backend/src/managers/modelsManager.spec.ts @@ -34,7 +34,7 @@ function mockFiles(now: Date) { const existsSyncSpy = vi.spyOn(fs, 'existsSync'); existsSyncSpy.mockImplementation((path: string) => { if (process.platform === 'win32') { - expect(path).toBe('\\home\\user\\aistudio\\models'); + expect(path).toBe('C:\\home\\user\\aistudio\\models'); } else { expect(path).toBe('/home/user/aistudio/models'); } @@ -62,7 +62,13 @@ function mockFiles(now: Date) { test('getLocalModelsFromDisk should get models in local directory', () => { const now = new Date(); mockFiles(now); - const manager = new ModelsManager('/home/user/aistudio', {} as Webview, {} as CatalogManager); + let appdir: string; + if (process.platform === 'win32') { + appdir = 'C:\\home\\user\\aistudio'; + } else { + appdir = '/home/user/aistudio'; + } + const manager = new ModelsManager(appdir, {} as Webview, {} as CatalogManager); manager.getLocalModelsFromDisk(); expect(manager.getLocalModels()).toEqual([ { @@ -86,11 +92,17 @@ test('getLocalModelsFromDisk should return an empty array if the models folder d vi.spyOn(os, 'homedir').mockReturnValue('/home/user'); const existsSyncSpy = vi.spyOn(fs, 'existsSync'); existsSyncSpy.mockReturnValue(false); - const manager = new ModelsManager('/home/user/aistudio', {} as Webview, {} as CatalogManager); + let appdir: string; + if (process.platform === 'win32') { + appdir = 'C:\\home\\user\\aistudio'; + } else { + appdir = '/home/user/aistudio'; + } + const manager = new ModelsManager(appdir, {} as Webview, {} as CatalogManager); manager.getLocalModelsFromDisk(); expect(manager.getLocalModels()).toEqual([]); if (process.platform === 'win32') { - expect(existsSyncSpy).toHaveBeenCalledWith('\\home\\user\\aistudio\\models'); + expect(existsSyncSpy).toHaveBeenCalledWith('C:\\home\\user\\aistudio\\models'); } else { expect(existsSyncSpy).toHaveBeenCalledWith('/home/user/aistudio/models'); } @@ -112,8 +124,14 @@ test('loadLocalModels should post a message with the message on disk and on cata }; }); const postMessageMock = vi.fn(); + let appdir: string; + if (process.platform === 'win32') { + appdir = 'C:\\home\\user\\aistudio'; + } else { + appdir = '/home/user/aistudio'; + } const manager = new ModelsManager( - '/home/user/aistudio', + appdir, { postMessage: postMessageMock, } as unknown as Webview, @@ -144,3 +162,134 @@ test('loadLocalModels should post a message with the message on disk and on cata ], }); }); + +test('deleteLocalModel deletes the model folder', async () => { + let appdir: string; + if (process.platform === 'win32') { + appdir = 'C:\\home\\user\\aistudio'; + } else { + appdir = '/home/user/aistudio'; + } + const now = new Date(); + mockFiles(now); + const rmSpy = vi.spyOn(fs.promises, 'rm'); + rmSpy.mockResolvedValue(); + const postMessageMock = vi.fn(); + const manager = new ModelsManager( + appdir, + { + postMessage: postMessageMock, + } as unknown as Webview, + { + getModels: () => { + return [ + { + id: 'model-id-1', + }, + ] as ModelInfo[]; + }, + } as CatalogManager, + ); + manager.getLocalModelsFromDisk(); + await manager.deleteLocalModel('model-id-1'); + // check that the model's folder is removed from disk + if (process.platform === 'win32') { + expect(rmSpy).toBeCalledWith('C:\\home\\user\\aistudio\\models\\model-id-1', { recursive: true }); + } else { + expect(rmSpy).toBeCalledWith('/home/user/aistudio/models/model-id-1', { recursive: true }); + } + expect(postMessageMock).toHaveBeenCalledTimes(2); + // check that a state is sent with the model being deleted + expect(postMessageMock).toHaveBeenCalledWith({ + id: 'new-local-models-state', + body: [ + { + file: { + creation: now, + file: 'model-id-1-model', + id: 'model-id-1', + size: 32000, + path: path.resolve(dirent[0].path, dirent[0].name, 'model-id-1-model'), + }, + id: 'model-id-1', + state: 'deleting', + }, + ], + }); + // check that a new state is sent with the model removed + expect(postMessageMock).toHaveBeenCalledWith({ + id: 'new-local-models-state', + body: [], + }); +}); + +test('deleteLocalModel fails to delete the model folder', async () => { + let appdir: string; + if (process.platform === 'win32') { + appdir = 'C:\\home\\user\\aistudio'; + } else { + appdir = '/home/user/aistudio'; + } + const now = new Date(); + mockFiles(now); + const rmSpy = vi.spyOn(fs.promises, 'rm'); + rmSpy.mockRejectedValue(new Error('failed')); + const postMessageMock = vi.fn(); + const manager = new ModelsManager( + appdir, + { + postMessage: postMessageMock, + } as unknown as Webview, + { + getModels: () => { + return [ + { + id: 'model-id-1', + }, + ] as ModelInfo[]; + }, + } as CatalogManager, + ); + manager.getLocalModelsFromDisk(); + await manager.deleteLocalModel('model-id-1'); + // check that the model's folder is removed from disk + if (process.platform === 'win32') { + expect(rmSpy).toBeCalledWith('C:\\home\\user\\aistudio\\models\\model-id-1', { recursive: true }); + } else { + expect(rmSpy).toBeCalledWith('/home/user/aistudio/models/model-id-1', { recursive: true }); + } + expect(postMessageMock).toHaveBeenCalledTimes(2); + // check that a state is sent with the model being deleted + expect(postMessageMock).toHaveBeenCalledWith({ + id: 'new-local-models-state', + body: [ + { + file: { + creation: now, + file: 'model-id-1-model', + id: 'model-id-1', + size: 32000, + path: path.resolve(dirent[0].path, dirent[0].name, 'model-id-1-model'), + }, + id: 'model-id-1', + state: 'deleting', + }, + ], + }); + // check that a new state is sent with the model non removed + expect(postMessageMock).toHaveBeenCalledWith({ + id: 'new-local-models-state', + body: [ + { + file: { + creation: now, + file: 'model-id-1-model', + id: 'model-id-1', + size: 32000, + path: path.resolve(dirent[0].path, dirent[0].name, 'model-id-1-model'), + }, + id: 'model-id-1', + }, + ], + }); +}); diff --git a/packages/backend/src/managers/modelsManager.ts b/packages/backend/src/managers/modelsManager.ts index c97a70f78..8a68f5e70 100644 --- a/packages/backend/src/managers/modelsManager.ts +++ b/packages/backend/src/managers/modelsManager.ts @@ -100,12 +100,16 @@ export class ModelsManager { return path.resolve(this.#modelsDir, modelId, info.file); } + getLocalModelFolder(modelId: string): string { + return path.resolve(this.#modelsDir, modelId); + } + getLocalModels(): LocalModelInfo[] { return Array.from(this.#localModels.values()); } async deleteLocalModel(modelId: string): Promise { - const modelDir = this.getLocalModelPath(modelId); + const modelDir = this.getLocalModelFolder(modelId); this.#deleted.add(modelId); await this.sendModelsInfo(); try {