fix: more tests

Signed-off-by: axel7083 <[email protected]>

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');
   });
 });

packages/backend/src/managers/models/UpdateManager.ts

@@ -62,7 +62,7 @@ export class UpdateManager extends Publisher<UpdateInfo[]> implements Disposable
 
       if (localEtag !== remoteEtag) {
         this.#updates.set(model.id, {
-          modelsId: model.id,
+          modelId: model.id,
           message: 'New update is available.',
         });
       }

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);
+  });
 });

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');
 

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();
+});

packages/frontend/src/lib/table/model/ModelColumnAction.spec.ts

@@ -27,22 +27,39 @@ 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', () => ({
   studioClient: {
     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');
+  });
+});

packages/frontend/src/lib/table/model/ModelColumnActions.svelte

@@ -1,6 +1,6 @@
 <script lang="ts">
 import type { ModelInfo } from '@shared/src/models/IModelInfo';
-import { faArrowUp, faCircleArrowUp, faDownload, faRocket, faTrash } from '@fortawesome/free-solid-svg-icons';
+import { faCircleArrowUp, faDownload, faRocket, faTrash } from '@fortawesome/free-solid-svg-icons';
 import { faFolderOpen } from '@fortawesome/free-solid-svg-icons';
 import ListItemButtonIcon from '../../button/ListItemButtonIcon.svelte';
 import { studioClient } from '/@/utils/client';
@@ -10,7 +10,7 @@ import type { UpdateInfo } from '@shared/src/models/IUpdate';
 export let object: ModelInfo;
 
 let updateInfo: UpdateInfo | undefined = undefined;
-$: updateInfo = $modelsUpdateInfo.find(update => update.modelsId === object.id);
+$: updateInfo = $modelsUpdateInfo.find(update => update.modelId === object.id);
 
 function deleteModel() {
   studioClient.requestRemoveLocalModel(object.id).catch(err => {
@@ -60,5 +60,5 @@ function updateModel() {
   <ListItemButtonIcon icon="{faDownload}" onClick="{downloadModel}" title="Download Model" enabled="{!object.state}" />
 {/if}
 {#if updateInfo}
-  <ListItemButtonIcon icon="{faCircleArrowUp}" title="Update model" onClick="{() => updateModel()}" />
+  <ListItemButtonIcon icon="{faCircleArrowUp}" title="{updateInfo.message}" onClick="{() => updateModel()}" />
 {/if}

packages/frontend/src/pages/Models.svelte

@@ -38,7 +38,7 @@ const columns: Column<ModelInfo>[] = [
     renderer: ModelColumnAge,
     comparator: (a, b) => (a.file?.creation?.getTime() ?? 0) - (b.file?.creation?.getTime() ?? 0),
   }),
-  new Column<ModelInfo>('', { width: '225px', align: 'right', renderer: ModelColumnLabels }),
+  new Column<ModelInfo>('', { width: '240px', align: 'right', renderer: ModelColumnLabels }),
   new Column<ModelInfo>('Actions', { align: 'right', width: '160px', renderer: ModelColumnActions }),
 ];
 const row = new Row<ModelInfo>({});

packages/shared/src/models/IUpdate.ts

@@ -18,5 +18,5 @@
 
 export interface UpdateInfo {
   message: string;
-  modelsId: string;
+  modelId: string;
 }