Skip to content

Commit

Permalink
feat(catalog): splitting logic in smaller utilities
Browse files Browse the repository at this point in the history 
Signed-off-by: axel7083 <[email protected]>
  • Loading branch information
@axel7083
axel7083 committed Mar 6, 2024
1 parent 2dc904d commit 1cd8a19
Show file tree
Hide file tree
Showing 6 changed files with 176 additions and 117 deletions.
48 changes: 27 additions & 21 deletions packages/backend/src/managers/catalogManager.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import { beforeEach, describe, expect, test, vi } from 'vitest';
import content from '../ai-test.json';
import userContent from '../ai-user-test.json';
import type { Webview } from '@podman-desktop/api';
import { EventEmitter, Webview } from '@podman-desktop/api';
import { CatalogManager } from '../managers/catalogManager';

import * as fs from 'node:fs';
Expand All @@ -41,24 +41,13 @@ vi.mock('node:fs', () => {
};
});

vi.mock('@podman-desktop/api', () => {
return {
fs: {
createFileSystemWatcher: () => ({
onDidCreate: vi.fn(),
onDidDelete: vi.fn(),
onDidChange: vi.fn(),
}),
},
};
});

const mocks = vi.hoisted(() => ({
withProgressMock: vi.fn(),
}));

vi.mock('@podman-desktop/api', async () => {
return {
EventEmitter: vi.fn(),
window: {
withProgress: mocks.withProgressMock,
},
Expand All @@ -78,20 +67,33 @@ vi.mock('@podman-desktop/api', async () => {
let catalogManager: CatalogManager;

beforeEach(async () => {
const appUserDirectory = '.';
vi.resetAllMocks();

const appUserDirectory = '.';
// Creating CatalogManager
catalogManager = new CatalogManager(appUserDirectory, {
postMessage: vi.fn(),
} as unknown as Webview);
vi.resetAllMocks();
catalogManager = new CatalogManager({
postMessage: vi.fn().mockResolvedValue(undefined),
} as unknown as Webview,
appUserDirectory, );

vi.mock('node:fs');

const listeners: ((value: unknown) => void)[] = [];

vi.mocked(EventEmitter).mockReturnValue({
event: vi.fn().mockImplementation((callback) => {
listeners.push(callback);
}),
fire: vi.fn().mockImplementation((content: unknown) => {
listeners.forEach((listener) => listener(content));
}),
} as unknown as EventEmitter<unknown>);
});

describe('invalid user catalog', () => {
beforeEach(async () => {
vi.spyOn(fs.promises, 'readFile').mockResolvedValue('invalid json');
await catalogManager.loadCatalog();
catalogManager.init();
});

test('expect correct model is returned with valid id', () => {
Expand All @@ -111,7 +113,9 @@ describe('invalid user catalog', () => {

test('expect correct model is returned from default catalog with valid id when no user catalog exists', async () => {
vi.spyOn(fs, 'existsSync').mockReturnValue(false);
await catalogManager.loadCatalog();
catalogManager.init();
await vi.waitUntil(() => catalogManager.getRecipes().length > 0);

const model = catalogManager.getModelById('hf.TheBloke.llama-2-7b-chat.Q5_K_S');
expect(model).toBeDefined();
expect(model.name).toEqual('TheBloke/Llama-2-7B-Chat-GGUF');
Expand All @@ -125,7 +129,9 @@ test('expect correct model is returned with valid id when the user catalog is va
vi.spyOn(fs, 'existsSync').mockReturnValue(true);
vi.spyOn(fs.promises, 'readFile').mockResolvedValue(JSON.stringify(userContent));

await catalogManager.loadCatalog();
catalogManager.init();
await vi.waitUntil(() => catalogManager.getRecipes().length > 0);

const model = catalogManager.getModelById('model1');
expect(model).toBeDefined();
expect(model.name).toEqual('Model 1');
Expand Down
101 changes: 26 additions & 75 deletions packages/backend/src/managers/catalogManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,42 +18,58 @@

import type { Catalog } from '@shared/src/models/ICatalog';
import path from 'node:path';
import { existsSync, promises } from 'node:fs';
import defaultCatalog from '../ai.json';
import type { Category } from '@shared/src/models/ICategory';
import type { Recipe } from '@shared/src/models/IRecipe';
import type { ModelInfo } from '@shared/src/models/IModelInfo';
import { MSG_NEW_CATALOG_STATE } from '@shared/Messages';
import { type Disposable, type Webview, fs } from '@podman-desktop/api';
import { JsonWatcher } from '../utils/JsonWatcher';
import { Publisher } from '../utils/Publisher';

export class CatalogManager implements Disposable {
private watchers: Map<string, Disposable> = new Map<string, Disposable>();
export class CatalogManager extends Publisher<Catalog> implements Disposable {
private catalog: Catalog;
#disposables: Disposable[];

constructor(
webview: Webview,
private appUserDirectory: string,
private webview: Webview,
) {
super(webview, MSG_NEW_CATALOG_STATE, () => this.getCatalog());
// We start with an empty catalog, for the methods to work before the catalog is loaded
this.catalog = {
categories: [],
models: [],
recipes: [],
};

this.#disposables = [];
}

init(): void {
// Creating a json watcher
const jsonWatcher: JsonWatcher<Catalog> = new JsonWatcher(
path.resolve(this.appUserDirectory, 'catalog.json'),
defaultCatalog
)
jsonWatcher.onContentUpdated((content) => this.onCatalogUpdated(content));
jsonWatcher.init();

this.#disposables.push(jsonWatcher);
}

private onCatalogUpdated(content: Catalog): void {
this.catalog = content;
this.notify();
}

dispose(): void {
Array.from(this.watchers.values()).forEach(watcher => watcher.dispose());
this.#disposables.forEach(watcher => watcher.dispose());
}

public getCatalog(): Catalog {
return this.catalog;
}

public getCategories(): Category[] {
return this.catalog.categories;
}

public getModels(): ModelInfo[] {
return this.catalog.models;
}
Expand All @@ -77,69 +93,4 @@ export class CatalogManager implements Disposable {
}
return recipe;
}

async loadCatalog() {
const catalogPath = path.resolve(this.appUserDirectory, 'catalog.json');

try {
this.watchCatalogFile(catalogPath); // do not await, we want to do this async
} catch (err: unknown) {
console.error(`unable to watch catalog file, changes to the catalog file won't be reflected to the UI`, err);
}

if (!existsSync(catalogPath)) {
return this.setCatalog(defaultCatalog);
}

try {
const cat = await this.readAndAnalyzeCatalog(catalogPath);
return this.setCatalog(cat);
} catch (err: unknown) {
console.error('unable to read catalog file, reverting to default catalog', err);
}
// If something went wrong we load the default catalog
return this.setCatalog(defaultCatalog);
}

watchCatalogFile(path: string) {
if (this.watchers.has(path)) throw new Error(`A watcher already exist for file ${path}.`);

const watcher = fs.createFileSystemWatcher(path);
this.watchers.set(path, watcher);

watcher.onDidCreate(async () => {
try {
const cat = await this.readAndAnalyzeCatalog(path);
await this.setCatalog(cat);
} catch (err: unknown) {
console.error('unable to read created catalog file, continue using default catalog', err);
}
});
watcher.onDidDelete(async () => {
console.log('user catalog file deleted, reverting to default catalog');
await this.setCatalog(defaultCatalog);
});
watcher.onDidChange(async () => {
try {
const cat = await this.readAndAnalyzeCatalog(path);
await this.setCatalog(cat);
} catch (err: unknown) {
console.error('unable to read modified catalog file, reverting to default catalog', err);
}
});
}

async readAndAnalyzeCatalog(path: string): Promise<Catalog> {
const data = await promises.readFile(path, 'utf-8');
return JSON.parse(data) as Catalog;
// TODO(feloy): check version, ...
}

async setCatalog(newCatalog: Catalog) {
this.catalog = newCatalog;
await this.webview.postMessage({
id: MSG_NEW_CATALOG_STATE,
body: this.catalog,
});
}
}
44 changes: 24 additions & 20 deletions packages/backend/src/studio-api-impl.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,15 @@ import userContent from './ai-user-test.json';
import type { ApplicationManager } from './managers/applicationManager';
import { StudioApiImpl } from './studio-api-impl';
import type { PlayGroundManager } from './managers/playground';
import type { TelemetryLogger, Webview } from '@podman-desktop/api';
import { EventEmitter, TelemetryLogger, Webview } from '@podman-desktop/api';
import { CatalogManager } from './managers/catalogManager';
import type { ModelsManager } from './managers/modelsManager';

import * as fs from 'node:fs';
import { timeout } from './utils/utils';
import type { TaskRegistry } from './registries/TaskRegistry';
import type { LocalRepositoryRegistry } from './registries/LocalRepositoryRegistry';
import { Recipe } from '@shared/src/models/IRecipe';

vi.mock('./ai.json', () => {
return {
Expand All @@ -48,18 +49,6 @@ vi.mock('node:fs', () => {
};
});

vi.mock('@podman-desktop/api', () => {
return {
fs: {
createFileSystemWatcher: () => ({
onDidCreate: vi.fn(),
onDidDelete: vi.fn(),
onDidChange: vi.fn(),
}),
},
};
});

const mocks = vi.hoisted(() => ({
withProgressMock: vi.fn(),
showWarningMessageMock: vi.fn(),
Expand All @@ -68,6 +57,7 @@ const mocks = vi.hoisted(() => ({

vi.mock('@podman-desktop/api', async () => {
return {
EventEmitter: vi.fn(),
window: {
withProgress: mocks.withProgressMock,
showWarningMessage: mocks.showWarningMessageMock,
Expand All @@ -86,15 +76,18 @@ vi.mock('@podman-desktop/api', async () => {
});

let studioApiImpl: StudioApiImpl;
let catalogManager;
let catalogManager: CatalogManager;

beforeEach(async () => {
vi.resetAllMocks();

const appUserDirectory = '.';

// Creating CatalogManager
catalogManager = new CatalogManager(appUserDirectory, {
postMessage: vi.fn(),
} as unknown as Webview);
catalogManager = new CatalogManager({
postMessage: vi.fn().mockResolvedValue(undefined),
} as unknown as Webview,
appUserDirectory);

// Creating StudioApiImpl
studioApiImpl = new StudioApiImpl(
Expand All @@ -108,8 +101,18 @@ beforeEach(async () => {
{} as LocalRepositoryRegistry,
{} as unknown as TaskRegistry,
);
vi.resetAllMocks();
vi.mock('node:fs');

const listeners: ((value: unknown) => void)[] = [];

vi.mocked(EventEmitter).mockReturnValue({
event: vi.fn().mockImplementation((callback) => {
listeners.push(callback);
}),
fire: vi.fn().mockImplementation((content: unknown) => {
listeners.forEach((listener) => listener(content));
}),
} as unknown as EventEmitter<unknown>);
});

test('expect pull application to call the withProgress api method', async () => {
Expand All @@ -118,15 +121,16 @@ test('expect pull application to call the withProgress api method', async () =>

mocks.withProgressMock.mockResolvedValue(undefined);

await catalogManager.loadCatalog();
catalogManager.init();
await vi.waitUntil(() => catalogManager.getRecipes().length > 0);
await studioApiImpl.pullApplication('recipe 1', 'model1');
expect(mocks.withProgressMock).toHaveBeenCalledOnce();
});

test('requestRemoveEnvironment should ask confirmation', async () => {
vi.spyOn(catalogManager, 'getRecipeById').mockReturnValue({
name: 'Recipe 1',
});
} as unknown as Recipe);
mocks.showWarningMessageMock.mockResolvedValue('Confirm');
await studioApiImpl.requestRemoveEnvironment('recipe-id-1', 'model-id-1');
await timeout(0);
Expand Down
2 changes: 1 addition & 1 deletion packages/backend/src/studio-api-impl.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ export class StudioApiImpl implements StudioAPI {

async pullApplication(recipeId: string, modelId: string): Promise<void> {
const recipe = this.catalogManager.getRecipes().find(recipe => recipe.id === recipeId);
if (!recipe) throw new Error('Not found');
if (!recipe) throw new Error(`recipe with if ${recipeId} not found`);

const model = this.catalogManager.getModelById(modelId);

Expand Down
Loading

0 comments on commit 1cd8a19

Please sign in to comment.