File size: 4,834 Bytes
4e23b01 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type {
IConfigService,
IProviderDiscoveryService,
ModelCatalogConfig,
} from '@moonshot-ai/agent-core-v2';
import { ModelCatalogRefreshScheduler } from '../src/services/modelCatalog/modelCatalogRefreshScheduler';
import type { ServerLogger } from '../src/services/pinoLoggerService';
const EMPTY_RESULT = { changed: [], unchanged: [], failed: [] };
function makeCatalog(refreshProviderModels = vi.fn(async () => EMPTY_RESULT)) {
return { refreshProviderModels } as unknown as IProviderDiscoveryService;
}
function makeConfig(catalogConfig?: ModelCatalogConfig): IConfigService {
return {
ready: Promise.resolve(),
get: vi.fn((domain: string) => (domain === 'modelCatalog' ? catalogConfig : undefined)),
} as unknown as IConfigService;
}
function makeLogger(): Pick<ServerLogger, 'info' | 'warn'> {
return {
info: vi.fn(),
warn: vi.fn(),
};
}
describe('ModelCatalogRefreshScheduler', () => {
beforeEach(() => {
vi.useFakeTimers();
});
afterEach(() => {
vi.useRealTimers();
vi.unstubAllEnvs();
});
it('refreshes on start and then on the configured interval', async () => {
const catalog = makeCatalog();
const scheduler = new ModelCatalogRefreshScheduler(catalog, makeConfig(), makeLogger(), {});
await scheduler.start();
await vi.waitFor(() => {
expect(catalog.refreshProviderModels).toHaveBeenCalledTimes(1);
});
expect(catalog.refreshProviderModels).toHaveBeenCalledWith({ scope: 'all' });
await vi.advanceTimersByTimeAsync(6 * 60 * 60 * 1000);
expect(catalog.refreshProviderModels).toHaveBeenCalledTimes(2);
scheduler.dispose();
await vi.advanceTimersByTimeAsync(6 * 60 * 60 * 1000);
expect(catalog.refreshProviderModels).toHaveBeenCalledTimes(2);
});
it('honors env overrides for interval and refresh-on-start', async () => {
const catalog = makeCatalog();
const scheduler = new ModelCatalogRefreshScheduler(catalog, makeConfig(), makeLogger(), {
KIMI_CODE_MODEL_CATALOG_REFRESH_INTERVAL_MS: '1000',
KIMI_CODE_MODEL_CATALOG_REFRESH_ON_START: '0',
});
await scheduler.start();
await vi.advanceTimersByTimeAsync(999);
expect(catalog.refreshProviderModels).not.toHaveBeenCalled();
await vi.advanceTimersByTimeAsync(1);
expect(catalog.refreshProviderModels).toHaveBeenCalledTimes(1);
});
it('reads interval and refresh-on-start from the modelCatalog config section', async () => {
const catalog = makeCatalog();
const scheduler = new ModelCatalogRefreshScheduler(
catalog,
makeConfig({ refreshIntervalMs: 1000, refreshOnStart: false }),
makeLogger(),
{},
);
await scheduler.start();
await vi.advanceTimersByTimeAsync(999);
expect(catalog.refreshProviderModels).not.toHaveBeenCalled();
await vi.advanceTimersByTimeAsync(1);
expect(catalog.refreshProviderModels).toHaveBeenCalledTimes(1);
});
it('lets env override the modelCatalog config section', async () => {
const catalog = makeCatalog();
const scheduler = new ModelCatalogRefreshScheduler(
catalog,
makeConfig({ refreshIntervalMs: 6 * 60 * 60 * 1000, refreshOnStart: true }),
makeLogger(),
{
KIMI_CODE_MODEL_CATALOG_REFRESH_ON_START: '0',
KIMI_CODE_MODEL_CATALOG_REFRESH_INTERVAL_MS: '1000',
},
);
await scheduler.start();
await vi.advanceTimersByTimeAsync(999);
expect(catalog.refreshProviderModels).not.toHaveBeenCalled();
await vi.advanceTimersByTimeAsync(1);
expect(catalog.refreshProviderModels).toHaveBeenCalledTimes(1);
});
it('disables the schedule when the config interval is 0', async () => {
const catalog = makeCatalog();
const scheduler = new ModelCatalogRefreshScheduler(
catalog,
makeConfig({ refreshIntervalMs: 0, refreshOnStart: false }),
makeLogger(),
{},
);
await scheduler.start();
await vi.advanceTimersByTimeAsync(60 * 1000);
expect(catalog.refreshProviderModels).not.toHaveBeenCalled();
});
it('swallows refresh errors so a failing tick does not break the schedule', async () => {
const refreshProviderModels = vi.fn().mockRejectedValue(new Error('network down'));
const catalog = makeCatalog(refreshProviderModels);
const logger = makeLogger();
const scheduler = new ModelCatalogRefreshScheduler(catalog, makeConfig(), logger, {
KIMI_CODE_MODEL_CATALOG_REFRESH_INTERVAL_MS: '1000',
KIMI_CODE_MODEL_CATALOG_REFRESH_ON_START: 'false',
});
await scheduler.start();
await vi.advanceTimersByTimeAsync(1000);
await vi.advanceTimersByTimeAsync(1000);
expect(refreshProviderModels).toHaveBeenCalledTimes(2);
expect(logger.warn).toHaveBeenCalled();
});
});
|