diff --git a/packages/backend/src/managers/models/UpdateManager.spec.ts b/packages/backend/src/managers/models/UpdateManager.spec.ts index 68a03376e..e2eaef78c 100644 --- a/packages/backend/src/managers/models/UpdateManager.spec.ts +++ b/packages/backend/src/managers/models/UpdateManager.spec.ts @@ -151,7 +151,7 @@ test('init call should ignore request without etag header', async () => { await vi.waitFor(() => { const updates = updater.getAll(); expect(updates.length).toBe(1); - expect(updates[0].modelsId).toBe('dummy-model-id-2'); + expect(updates[0].modelId).toBe('dummy-model-id-2'); }); }); @@ -194,7 +194,7 @@ test('models info without file should not be checked', async () => { await vi.waitFor(() => { const updates = updater.getAll(); expect(updates.length).toBe(1); - expect(updates[0].modelsId).toBe('dummy-model-id-2'); + expect(updates[0].modelId).toBe('dummy-model-id-2'); }); }); @@ -243,6 +243,6 @@ test('models info without local etag should not be checked', async () => { await vi.waitFor(() => { const updates = updater.getAll(); expect(updates.length).toBe(1); - expect(updates[0].modelsId).toBe('dummy-model-id-2'); + expect(updates[0].modelId).toBe('dummy-model-id-2'); }); }); diff --git a/packages/backend/src/managers/models/UpdateManager.ts b/packages/backend/src/managers/models/UpdateManager.ts index 8ea3b812c..2d4ef1c3b 100644 --- a/packages/backend/src/managers/models/UpdateManager.ts +++ b/packages/backend/src/managers/models/UpdateManager.ts @@ -62,7 +62,7 @@ export class UpdateManager extends Publisher implements Disposable if (localEtag !== remoteEtag) { this.#updates.set(model.id, { - modelsId: model.id, + modelId: model.id, message: 'New update is available.', }); } diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts index a24f22046..e36779aa6 100644 --- a/packages/backend/src/managers/modelsManager.spec.ts +++ b/packages/backend/src/managers/modelsManager.spec.ts @@ -774,4 +774,33 @@ describe('downloadModel', () => { expect(mocks.performDownloadMock).toHaveBeenCalledTimes(1); expect(mocks.onEventDownloadMock).toHaveBeenCalledTimes(2); }); + + test('using is-update label should force models to be re-downloaded', async () => { + const manager = new ModelsManager( + 'appdir', + {} as Webview, + { + getModels(): ModelInfo[] { + return []; + }, + } as CatalogManager, + telemetryLogger, + taskRegistry, + cancellationTokenRegistryMock, + ); + vi.spyOn(manager, 'isModelOnDisk').mockReturnValue(true); + await manager.requestDownloadModel( + { + id: 'id', + url: 'url', + name: 'name', + } as ModelInfo, + { + 'is-update': 'true', + }, + ); + + expect(mocks.performDownloadMock).toHaveBeenCalledTimes(1); + expect(mocks.onEventDownloadMock).toHaveBeenCalledTimes(1); + }); }); diff --git a/packages/backend/src/studio-api-impl.spec.ts b/packages/backend/src/studio-api-impl.spec.ts index 3ea73130c..3b5bf8819 100644 --- a/packages/backend/src/studio-api-impl.spec.ts +++ b/packages/backend/src/studio-api-impl.spec.ts @@ -38,6 +38,7 @@ import type { CancellationTokenRegistry } from './registries/CancellationTokenRe import path from 'node:path'; import type { LocalModelImportInfo } from '@shared/src/models/ILocalModelInfo'; import * as podman from './utils/podman'; +import type { UpdateManager } from './managers/models/UpdateManager'; vi.mock('./ai.json', () => { return { @@ -138,6 +139,7 @@ beforeEach(async () => { {} as unknown as PlaygroundV2Manager, {} as unknown as SnippetManager, {} as unknown as CancellationTokenRegistry, + {} as unknown as UpdateManager, ); vi.mock('node:fs'); diff --git a/packages/backend/src/utils/downloader.spec.ts b/packages/backend/src/utils/downloader.spec.ts index 2702ed622..48ee4af16 100644 --- a/packages/backend/src/utils/downloader.spec.ts +++ b/packages/backend/src/utils/downloader.spec.ts @@ -96,8 +96,11 @@ test('perform download failed', async () => { const listenerMock = vi.fn(); downloader.onEvent(listenerMock); + let aborted: boolean | undefined = undefined; // perform download logic (do not wait) - void downloader.perform('followUpId'); + void downloader.perform('followUpId').then(result => { + aborted = result; + }); // wait for listener to be registered await vi.waitFor(() => { @@ -116,6 +119,8 @@ test('perform download failed', async () => { expect(downloader.completed).toBeTruthy(); }); + expect(aborted).toBe(false); + expect(listenerMock).toHaveBeenCalledWith({ id: 'followUpId', message: 'Something went wrong: dummyError.', @@ -179,3 +184,52 @@ test('perform download successfully', async () => { }); expect(promises.rm).not.toHaveBeenCalled(); }); + +test('aborted signal should return true to perform', async () => { + const controller = new AbortController(); + const downloader = new Downloader('dummyUrl', 'dummyTarget', controller.signal); + vi.mocked(https.get).mockImplementation((_url, _options, callback) => { + callback?.({ + pipe: vi.fn(), + on: vi.fn(), + headers: { location: undefined }, + } as unknown as IncomingMessage); + return {} as unknown as ClientRequest; + }); + + const closeMock = vi.fn(); + const onMock = vi.fn(); + vi.mocked(createWriteStream).mockReturnValue({ + close: closeMock, + on: onMock, + } as unknown as WriteStream); + + onMock.mockImplementation((event: string, callback: () => void) => { + if (event === 'finish') { + callback(); + } + }); + + // capture downloader event(s) + const listenerMock = vi.fn(); + downloader.onEvent(listenerMock); + + // abort controller + controller.abort('testing'); + + // perform download logic + const aborted = await downloader.perform('followUpId'); + + expect(downloader.completed).toBeTruthy(); + expect(aborted).toBeTruthy(); + + expect(promises.rename).toHaveBeenCalledWith('dummyTarget.tmp', 'dummyTarget'); + expect(downloader.completed).toBeTruthy(); + expect(listenerMock).toHaveBeenCalledWith({ + id: 'followUpId', + duration: expect.anything(), + message: expect.anything(), + status: 'completed', + }); + expect(promises.rm).not.toHaveBeenCalled(); +}); diff --git a/packages/frontend/src/lib/table/model/ModelColumnAction.spec.ts b/packages/frontend/src/lib/table/model/ModelColumnAction.spec.ts index fed542234..e26fc3755 100644 --- a/packages/frontend/src/lib/table/model/ModelColumnAction.spec.ts +++ b/packages/frontend/src/lib/table/model/ModelColumnAction.spec.ts @@ -27,6 +27,17 @@ const mocks = vi.hoisted(() => ({ requestRemoveLocalModel: vi.fn(), openFile: vi.fn(), downloadModel: vi.fn(), + requestModelUpdate: vi.fn(), + getModelsUpdateInfoMock: vi.fn(), +})); + +vi.mock('/@/stores/modelsUpdateInfo', () => ({ + modelsUpdateInfo: { + subscribe: (f: (msg: any) => void) => { + f(mocks.getModelsUpdateInfoMock()); + return () => {}; + }, + }, })); vi.mock('/@/utils/client', () => ({ @@ -34,15 +45,21 @@ vi.mock('/@/utils/client', () => ({ requestRemoveLocalModel: mocks.requestRemoveLocalModel, openFile: mocks.openFile, downloadModel: mocks.downloadModel, + requestModelUpdate: mocks.requestModelUpdate, }, })); beforeEach(() => { vi.resetAllMocks(); + // mock store content + mocks.getModelsUpdateInfoMock.mockReturnValue([]); + + // mocks studio client methods mocks.downloadModel.mockResolvedValue(undefined); mocks.openFile.mockResolvedValue(undefined); mocks.requestRemoveLocalModel.mockResolvedValue(undefined); + mocks.requestModelUpdate.mockResolvedValue(undefined); }); test('Expect folder and delete button in document', async () => { @@ -160,3 +177,38 @@ test('Expect router to be called when rocket icon clicked', async () => { expect(replaceMock).toHaveBeenCalledWith({ 'model-id': 'my-model' }); }); }); + +test('Expect update icon to be visible when one available', async () => { + mocks.getModelsUpdateInfoMock.mockReturnValue([ + { + modelId: 'my-model', + message: 'New version available', + }, + ]); + + const object: ModelInfo = { + id: 'my-model', + description: '', + hw: '', + license: '', + name: '', + registry: '', + url: '', + file: { + file: 'file', + creation: new Date(), + size: 1000, + path: 'path', + }, + memory: 1000, + }; + render(ModelColumnActions, { object }); + + const updateBtn = screen.getByTitle('New version available'); + expect(updateBtn).toBeDefined(); + + await fireEvent.click(updateBtn); + await waitFor(() => { + expect(mocks.requestModelUpdate).toHaveBeenCalledWith('my-model'); + }); +}); diff --git a/packages/frontend/src/lib/table/model/ModelColumnActions.svelte b/packages/frontend/src/lib/table/model/ModelColumnActions.svelte index dcf23c5d7..2d3ddab2a 100644 --- a/packages/frontend/src/lib/table/model/ModelColumnActions.svelte +++ b/packages/frontend/src/lib/table/model/ModelColumnActions.svelte @@ -1,6 +1,6 @@