diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts index 47f1faf55..4cc4ac164 100644 --- a/packages/backend/src/managers/modelsManager.spec.ts +++ b/packages/backend/src/managers/modelsManager.spec.ts @@ -144,3 +144,114 @@ test('loadLocalModels should post a message with the message on disk and on cata ], }); }); + +test('deleteLocalModel deletes the model folder', async () => { + const now = new Date(); + mockFiles(now); + const rmSpy = vi.spyOn(fs.promises, 'rm'); + rmSpy.mockResolvedValue(); + const postMessageMock = vi.fn(); + const manager = new ModelsManager( + '/home/user/aistudio', + { + 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 + 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 () => { + 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( + '/home/user/aistudio', + { + 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 + 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 {