import { setupDatabaseAndSynchronizer, switchClient, createTempDir } from '../../testing/test-utils'; import * as fs from 'fs-extra'; import * as tar from 'tar'; import Setting from '../../models/Setting'; import shim from '../../shim'; import { ensureModelDownloaded, localModelPath, removeCachedModel, ModelDescriptor, } from './EmbeddingModelDownloader'; // The downloader test never touches the real network. We build a tiny // model-shaped tarball on disk and stub shim.fetchBlob to "download" it by // copying that tarball into the requested destination. Everything past that // (extract, marker check, cleanup) runs against the real fs. describe('EmbeddingModelDownloader', () => { let tempDir: string; let tarballPath: string; let originalFetchBlob: typeof shim.fetchBlob; const fakeModel: ModelDescriptor = { id: 'fake-test-model', archiveName: 'fake-test-model', downloadUrl: 'https://example.invalid/fake-test-model.tar.gz', }; const buildFakeModelTarball = async (): Promise => { // Lay out the staging dir to match the real archive: files at the // archive root, no top-level wrapper directory. The downloader // extracts into modelDir(model) directly. const stagingDir = `${tempDir}/staging`; await fs.ensureDir(stagingDir); await fs.writeFile(`${stagingDir}/config.json`, '{"_fake":true}'); await fs.writeFile(`${stagingDir}/tokenizer.json`, '{}'); const archive = `${tempDir}/${fakeModel.archiveName}.tar.gz`; await tar.create({ gzip: true, cwd: stagingDir, file: archive }, ['config.json', 'tokenizer.json']); return archive; }; beforeEach(async () => { await setupDatabaseAndSynchronizer(1); await switchClient(1); tempDir = await createTempDir(); // Point cacheDir at our scratch dir so the downloader can't pollute // other tests' cache state. Setting.setConstant('cacheDir', tempDir); tarballPath = await buildFakeModelTarball(); // Stub fetchBlob: instead of going to the network, copy our local // fake tarball to the requested destination. originalFetchBlob = shim.fetchBlob; // eslint-disable-next-line @typescript-eslint/no-explicit-any -- fetchBlob's options bag is loosely typed across implementations shim.fetchBlob = (async (_url: string, options: any) => { await fs.copy(tarballPath, options.path); return { ok: true, status: 200, path: options.path }; }); }); afterEach(async () => { shim.fetchBlob = originalFetchBlob; await fs.remove(tempDir); }); it('returns null from localModelPath when nothing is cached', async () => { expect(await localModelPath(fakeModel)).toBeNull(); }); it('downloads, extracts, and exposes the model directory', async () => { const path = await ensureModelDownloaded(fakeModel); expect(path).toMatch(/fake-test-model$/); expect(await fs.pathExists(`${path}/config.json`)).toBe(true); expect(await fs.pathExists(`${path}/tokenizer.json`)).toBe(true); // Tarball cleanup happens after successful extract. expect(await fs.pathExists(`${tempDir}/ai/embedding-models/${fakeModel.archiveName}.tar.gz`)).toBe(false); }); it('is idempotent: a second call returns the same path without re-downloading', async () => { await ensureModelDownloaded(fakeModel); // Count fetchBlob calls in a second run. let calls = 0; shim.fetchBlob = (async () => { calls++; return { ok: true, status: 200 }; }); const path = await ensureModelDownloaded(fakeModel); expect(calls).toBe(0); expect(await fs.pathExists(`${path}/config.json`)).toBe(true); }); it('serialises concurrent calls so only one download runs', async () => { let calls = 0; // Slow the "download" down to widen the race window — without // concurrency control, both calls would race to wipe + write the // same files at the same time. // eslint-disable-next-line @typescript-eslint/no-explicit-any -- fetchBlob spy mirrors the loose typing of its callers shim.fetchBlob = (async (_url: string, options: any) => { calls++; await new Promise(resolve => setTimeout(resolve, 50)); await fs.copy(tarballPath, options.path); return { ok: true, status: 200 }; // eslint-disable-next-line @typescript-eslint/no-explicit-any -- See above }) as any; const [p1, p2, p3] = await Promise.all([ ensureModelDownloaded(fakeModel), ensureModelDownloaded(fakeModel), ensureModelDownloaded(fakeModel), ]); expect(calls).toBe(1); expect(p1).toBe(p2); expect(p2).toBe(p3); expect(await fs.pathExists(`${p1}/config.json`)).toBe(true); }); it('removes any stale partial state before re-downloading', async () => { // Simulate a partial extract left over from a previous crash: directory // exists but is missing the marker file. const targetDir = `${tempDir}/ai/embedding-models/${fakeModel.archiveName}`; await fs.ensureDir(targetDir); await fs.writeFile(`${targetDir}/leftover.bin`, 'garbage'); expect(await localModelPath(fakeModel)).toBeNull(); const path = await ensureModelDownloaded(fakeModel); expect(await fs.pathExists(`${path}/config.json`)).toBe(true); // Leftover from the prior failed attempt must be gone. expect(await fs.pathExists(`${targetDir}/leftover.bin`)).toBe(false); }); it('throws a clear error when the server returns a non-2xx response', async () => { shim.fetchBlob = (async () => ({ ok: false, status: 404 })); await expect(ensureModelDownloaded(fakeModel)).rejects.toThrow(/HTTP 404/); }); it('removeCachedModel deletes the extracted directory', async () => { await ensureModelDownloaded(fakeModel); expect(await localModelPath(fakeModel)).not.toBeNull(); await removeCachedModel(fakeModel); expect(await localModelPath(fakeModel)).toBeNull(); }); it('reports progress via the onProgress callback', async () => { // Re-stub fetchBlob to actually exercise downloadController so we // observe chunk callbacks. We simulate a streaming download by // invoking handleChunk a few times before copying the tarball. // eslint-disable-next-line @typescript-eslint/no-explicit-any -- See above shim.fetchBlob = (async (_url: string, options: any) => { const ctrl = options.downloadController; if (ctrl?.handleChunk) { const onData = ctrl.handleChunk({ destroy: () => {} }); onData({ length: 1000 }); onData({ length: 500 }); } await fs.copy(tarballPath, options.path); return { ok: true, status: 200 }; }); const samples: { bytesDownloaded: number }[] = []; await ensureModelDownloaded(fakeModel, { onProgress: (p) => samples.push({ bytesDownloaded: p.bytesDownloaded }), }); expect(samples.length).toBeGreaterThanOrEqual(2); expect(samples[samples.length - 1].bytesDownloaded).toBe(1500); }); });