Update Astra: v2.80.19 - Refactoring Sidebar, LM Studio integration, and new tests
This commit is contained in:
g1nation
@@ -0,0 +1,310 @@
|
||||
/**
|
||||
* Unit tests for ModelLifecycleManager.
|
||||
*
|
||||
* Strategy: inject mock ILMStudioClient and a simple in-memory IActivityTracker.
|
||||
* No real LM Studio or SDK is touched — the manager file does not import the
|
||||
* SDK directly (only types via `import type`).
|
||||
*/
|
||||
|
||||
import {
|
||||
ModelLifecycleManager,
|
||||
LifecycleConfig,
|
||||
LifecycleManagerDeps,
|
||||
} from '../src/lmstudio/lifecycleManager';
|
||||
import type { ILMStudioClient } from '../src/lmstudio/client';
|
||||
import type { IActivityTracker } from '../src/lmstudio/activityTracker';
|
||||
|
||||
class FakeActivityTracker implements IActivityTracker {
|
||||
private listeners: Array<(_: void) => any> = [];
|
||||
public readonly onActivity = ((listener: (_: void) => any) => {
|
||||
this.listeners.push(listener);
|
||||
return { dispose: () => { this.listeners = this.listeners.filter(l => l !== listener); } };
|
||||
}) as any;
|
||||
bump(): void {
|
||||
for (const l of this.listeners.slice()) l();
|
||||
}
|
||||
}
|
||||
|
||||
class FakeLMStudioClient implements ILMStudioClient {
|
||||
public loadCalls: string[] = [];
|
||||
public unloadCalls: string[] = [];
|
||||
public listLoadedCalls = 0;
|
||||
public failNextLoad: Error | null = null;
|
||||
public failNextUnload: Error | null = null;
|
||||
public loadDelayMs = 0;
|
||||
public lastLoadSignal: AbortSignal | undefined;
|
||||
|
||||
setBaseUrl(_: string): void { /* noop */ }
|
||||
|
||||
async load(modelKey: string, signal?: AbortSignal): Promise<void> {
|
||||
this.loadCalls.push(modelKey);
|
||||
this.lastLoadSignal = signal;
|
||||
if (this.loadDelayMs > 0) {
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
const t = setTimeout(resolve, this.loadDelayMs);
|
||||
if (signal) {
|
||||
const onAbort = () => { clearTimeout(t); reject(new Error('aborted')); };
|
||||
if (signal.aborted) onAbort();
|
||||
else signal.addEventListener('abort', onAbort);
|
||||
}
|
||||
});
|
||||
}
|
||||
if (this.failNextLoad) {
|
||||
const err = this.failNextLoad;
|
||||
this.failNextLoad = null;
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
async unload(modelKey: string): Promise<void> {
|
||||
this.unloadCalls.push(modelKey);
|
||||
if (this.failNextUnload) {
|
||||
const err = this.failNextUnload;
|
||||
this.failNextUnload = null;
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
async listLoaded(): Promise<string[]> {
|
||||
this.listLoadedCalls++;
|
||||
return [];
|
||||
}
|
||||
|
||||
async isReachable(): Promise<boolean> {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
function makeManager(overrides: Partial<LifecycleConfig> = {}, depOverrides: Partial<LifecycleManagerDeps> = {}) {
|
||||
const client = new FakeLMStudioClient();
|
||||
const activity = new FakeActivityTracker();
|
||||
const config: LifecycleConfig = { idleTimeoutMs: 1000, autoLoadOnSelect: true, ...overrides };
|
||||
const errors: string[] = [];
|
||||
const manager = new ModelLifecycleManager({
|
||||
client,
|
||||
activity,
|
||||
getConfig: () => config,
|
||||
notifyError: (m) => errors.push(m),
|
||||
switchDebounceMs: 0,
|
||||
initialEngine: 'lmstudio',
|
||||
...depOverrides,
|
||||
});
|
||||
return { manager, client, activity, config, errors };
|
||||
}
|
||||
|
||||
const flush = () => new Promise<void>((r) => setImmediate(r));
|
||||
|
||||
describe('ModelLifecycleManager', () => {
|
||||
beforeEach(() => {
|
||||
jest.useFakeTimers({ doNotFake: ['setImmediate'] });
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
jest.useRealTimers();
|
||||
});
|
||||
|
||||
test('loads model on selection and arms idle timer', async () => {
|
||||
const { manager, client } = makeManager();
|
||||
manager.onModelSelected('llama-3.2-3b');
|
||||
await flush();
|
||||
expect(client.loadCalls).toEqual(['llama-3.2-3b']);
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
expect(manager._getCurrentModel()).toBe('llama-3.2-3b');
|
||||
expect(manager._hasIdleTimer()).toBe(true);
|
||||
});
|
||||
|
||||
test('idle timer triggers unload after timeout', async () => {
|
||||
const { manager, client } = makeManager({ idleTimeoutMs: 1000 });
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
|
||||
jest.advanceTimersByTime(1000);
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual(['m1']);
|
||||
expect(manager._getState()).toBe('idle');
|
||||
expect(manager._getCurrentModel()).toBe(null);
|
||||
});
|
||||
|
||||
test('idleTimeoutMs <= 0 disables auto-eject', async () => {
|
||||
const { manager, client } = makeManager({ idleTimeoutMs: 0 });
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
expect(manager._hasIdleTimer()).toBe(false);
|
||||
|
||||
jest.advanceTimersByTime(1_000_000);
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual([]);
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
});
|
||||
|
||||
test('activity bump resets idle timer', async () => {
|
||||
const { manager, client, activity } = makeManager({ idleTimeoutMs: 1000 });
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
|
||||
jest.advanceTimersByTime(900);
|
||||
activity.bump();
|
||||
jest.advanceTimersByTime(900); // total 1800ms but timer reset at 900
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual([]);
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
|
||||
jest.advanceTimersByTime(200); // 200 + 900 since reset = ~1100 since reset
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual(['m1']);
|
||||
});
|
||||
|
||||
test('streamStart pauses idle timer; streamEnd resumes', async () => {
|
||||
const { manager, client } = makeManager({ idleTimeoutMs: 500 });
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(manager._hasIdleTimer()).toBe(true);
|
||||
|
||||
manager.onStreamStart();
|
||||
expect(manager._getState()).toBe('streaming');
|
||||
expect(manager._hasIdleTimer()).toBe(false);
|
||||
|
||||
jest.advanceTimersByTime(10000);
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual([]); // never ejected during stream
|
||||
|
||||
manager.onStreamEnd();
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
expect(manager._hasIdleTimer()).toBe(true);
|
||||
|
||||
jest.advanceTimersByTime(500);
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual(['m1']);
|
||||
});
|
||||
|
||||
test('model switch unloads previous and loads next', async () => {
|
||||
const { manager, client } = makeManager();
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
|
||||
manager.onModelSelected('m2');
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual(['m1']);
|
||||
expect(client.loadCalls).toEqual(['m1', 'm2']);
|
||||
expect(manager._getCurrentModel()).toBe('m2');
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
});
|
||||
|
||||
test('switch during streaming defers until streamEnd', async () => {
|
||||
const { manager, client } = makeManager();
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
manager.onStreamStart();
|
||||
|
||||
manager.onModelSelected('m2');
|
||||
await flush();
|
||||
// No switch yet — still in streaming with m1
|
||||
expect(client.unloadCalls).toEqual([]);
|
||||
expect(client.loadCalls).toEqual(['m1']);
|
||||
|
||||
manager.onStreamEnd();
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual(['m1']);
|
||||
expect(client.loadCalls).toEqual(['m1', 'm2']);
|
||||
expect(manager._getCurrentModel()).toBe('m2');
|
||||
});
|
||||
|
||||
test('load failure surfaces via notifyError and resets to idle', async () => {
|
||||
const { manager, client, errors } = makeManager();
|
||||
client.failNextLoad = new Error('LM Studio is not running');
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(manager._getState()).toBe('idle');
|
||||
expect(manager._getCurrentModel()).toBe(null);
|
||||
expect(errors.length).toBe(1);
|
||||
expect(errors[0]).toContain('LM Studio is not running');
|
||||
});
|
||||
|
||||
test('engine change to non-lmstudio clears state without unloading', async () => {
|
||||
const { manager, client } = makeManager();
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(manager._getState()).toBe('loaded');
|
||||
|
||||
manager.setEngine('ollama');
|
||||
expect(manager._getState()).toBe('idle');
|
||||
expect(manager._getCurrentModel()).toBe(null);
|
||||
expect(manager._hasIdleTimer()).toBe(false);
|
||||
expect(client.unloadCalls).toEqual([]); // explicitly does not call unload
|
||||
});
|
||||
|
||||
test('selection while engine is non-lmstudio is a no-op', async () => {
|
||||
const { manager, client } = makeManager({}, { initialEngine: 'ollama' });
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(client.loadCalls).toEqual([]);
|
||||
expect(manager._getState()).toBe('idle');
|
||||
});
|
||||
|
||||
test('autoLoadOnSelect=false skips loading', async () => {
|
||||
const { manager, client } = makeManager({ autoLoadOnSelect: false });
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(client.loadCalls).toEqual([]);
|
||||
expect(manager._getState()).toBe('idle');
|
||||
});
|
||||
|
||||
test('rapid switch debounce: only the last selection wins', async () => {
|
||||
// Re-create manager with non-zero debounce
|
||||
const client = new FakeLMStudioClient();
|
||||
const activity = new FakeActivityTracker();
|
||||
const config: LifecycleConfig = { idleTimeoutMs: 1000, autoLoadOnSelect: true };
|
||||
const manager = new ModelLifecycleManager({
|
||||
client,
|
||||
activity,
|
||||
getConfig: () => config,
|
||||
switchDebounceMs: 300,
|
||||
initialEngine: 'lmstudio',
|
||||
});
|
||||
manager.onModelSelected('m1');
|
||||
manager.onModelSelected('m2');
|
||||
manager.onModelSelected('m3');
|
||||
// Before debounce expires no load fires
|
||||
expect(client.loadCalls).toEqual([]);
|
||||
jest.advanceTimersByTime(300);
|
||||
await flush();
|
||||
expect(client.loadCalls).toEqual(['m3']);
|
||||
});
|
||||
|
||||
test('disposeAndUnload ejects loaded model', async () => {
|
||||
const { manager, client } = makeManager();
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
await manager.disposeAndUnload(2000);
|
||||
expect(client.unloadCalls).toEqual(['m1']);
|
||||
expect(manager._getState()).toBe('idle');
|
||||
});
|
||||
|
||||
test('disposeAndUnload while idle is a no-op', async () => {
|
||||
const { manager, client } = makeManager();
|
||||
await manager.disposeAndUnload(500);
|
||||
expect(client.unloadCalls).toEqual([]);
|
||||
});
|
||||
|
||||
test('selecting same model does not re-load but refreshes timer', async () => {
|
||||
const { manager, client } = makeManager({ idleTimeoutMs: 1000 });
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(client.loadCalls.length).toBe(1);
|
||||
|
||||
jest.advanceTimersByTime(800);
|
||||
manager.onModelSelected('m1');
|
||||
await flush();
|
||||
expect(client.loadCalls.length).toBe(1); // no reload
|
||||
// Timer was reset; should survive past original 1000ms
|
||||
jest.advanceTimersByTime(700);
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual([]);
|
||||
|
||||
jest.advanceTimersByTime(400);
|
||||
await flush();
|
||||
expect(client.unloadCalls).toEqual(['m1']);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user