add backend unit tests

containers · Jan 25, 2024 · 8e3d2ce · 8e3d2ce
1 parent bd11b63
commit 8e3d2ce
Show file tree

Hide file tree

Showing 2 changed files with 136 additions and 1 deletion.
diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts
@@ -144,3 +144,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('D:\\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
@@ -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<void> {
-    const modelDir = this.getLocalModelPath(modelId);
+    const modelDir = this.getLocalModelFolder(modelId);
     this.#deleted.add(modelId);
     await this.sendModelsInfo();
     try {