Skip to content

Commit

Permalink
add backend unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@feloy
feloy committed Jan 25, 2024
1 parent bd11b63 commit 408a6d8
Show file tree
Hide file tree
Showing 2 changed files with 159 additions and 6 deletions.
159 changes: 154 additions & 5 deletions packages/backend/src/managers/modelsManager.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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');
}
Expand Down Expand Up @@ -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([
{
Expand All @@ -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');
}
Expand All @@ -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,
Expand Down Expand Up @@ -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',
},
],
});
});
6 changes: 5 additions & 1 deletion packages/backend/src/managers/modelsManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 408a6d8

Please sign in to comment.