Spaces:
Running
Running
File size: 4,165 Bytes
7104219 | 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 | import fs from "node:fs";
import os from "node:os";
import path from "node:path";
import { downloadFile, fileDownloadInfo } from "@huggingface/hub";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { downloadFileFromHuggingFaceRepository } from "./downloadFileFromHuggingFaceRepository";
vi.mock("@huggingface/hub", () => ({
downloadFile: vi.fn(),
fileDownloadInfo: vi.fn(),
}));
const REPO = "jinaai/jina-reranker-v1-tiny-en";
const REPO_FILE = "onnx/model.onnx";
const REMOTE_CONTENT = "complete model bytes";
let temporaryDirectory: string;
let localFilePath: string;
function serveRemoteFile(content = REMOTE_CONTENT) {
vi.mocked(fileDownloadInfo).mockResolvedValue({
size: Buffer.byteLength(REMOTE_CONTENT),
etag: "etag",
url: `https://huggingface.co/${REPO}/resolve/main/${REPO_FILE}`,
});
vi.mocked(downloadFile).mockImplementation(async () => new Blob([content]));
}
function download(filePath = localFilePath) {
return downloadFileFromHuggingFaceRepository(REPO, REPO_FILE, filePath);
}
function listDirectory(directory = path.dirname(localFilePath)) {
return fs.existsSync(directory) ? fs.readdirSync(directory).sort() : [];
}
beforeEach(() => {
vi.clearAllMocks();
temporaryDirectory = fs.mkdtempSync(path.join(os.tmpdir(), "hf-download-"));
localFilePath = path.join(temporaryDirectory, "onnx", "model.onnx");
});
afterEach(() => {
vi.restoreAllMocks();
fs.rmSync(temporaryDirectory, { recursive: true, force: true });
});
describe("downloadFileFromHuggingFaceRepository", () => {
it("downloads a file that is not cached yet", async () => {
serveRemoteFile();
await download();
expect(fs.readFileSync(localFilePath, "utf8")).toBe(REMOTE_CONTENT);
expect(listDirectory()).toEqual(["model.onnx"]);
});
it("keeps a cached file whose size matches the repository", async () => {
serveRemoteFile();
fs.mkdirSync(path.dirname(localFilePath), { recursive: true });
fs.writeFileSync(localFilePath, REMOTE_CONTENT);
await download();
expect(downloadFile).not.toHaveBeenCalled();
});
it("only downloads the missing file when a sibling is already cached", async () => {
serveRemoteFile();
const cachedFilePath = path.join(temporaryDirectory, "tokenizer.json");
fs.writeFileSync(cachedFilePath, REMOTE_CONTENT);
await download(cachedFilePath);
await download();
expect(downloadFile).toHaveBeenCalledTimes(1);
expect(fs.readFileSync(localFilePath, "utf8")).toBe(REMOTE_CONTENT);
});
it("replaces a truncated cached file instead of trusting it", async () => {
serveRemoteFile();
fs.mkdirSync(path.dirname(localFilePath), { recursive: true });
fs.writeFileSync(localFilePath, REMOTE_CONTENT.slice(0, 5));
await download();
expect(downloadFile).toHaveBeenCalledTimes(1);
expect(fs.readFileSync(localFilePath, "utf8")).toBe(REMOTE_CONTENT);
});
it("keeps a cached file when the repository metadata is unreachable", async () => {
vi.mocked(fileDownloadInfo).mockRejectedValue(new Error("offline"));
fs.mkdirSync(path.dirname(localFilePath), { recursive: true });
fs.writeFileSync(localFilePath, REMOTE_CONTENT.slice(0, 5));
await download();
expect(downloadFile).not.toHaveBeenCalled();
expect(fs.readFileSync(localFilePath, "utf8")).toBe(
REMOTE_CONTENT.slice(0, 5),
);
});
it("writes nothing when the response is shorter than the expected size", async () => {
serveRemoteFile("truncated");
await expect(download()).rejects.toThrow(/9 bytes instead of 20/);
expect(listDirectory()).toEqual([]);
});
it("leaves no partial file behind when the disk fills up mid-write", async () => {
serveRemoteFile();
vi.spyOn(fs, "writeFileSync").mockImplementation((filePath) => {
const fileDescriptor = fs.openSync(filePath as string, "w");
fs.writeSync(fileDescriptor, REMOTE_CONTENT.slice(0, 5));
fs.closeSync(fileDescriptor);
throw new Error("ENOSPC: no space left on device");
});
await expect(download()).rejects.toThrow("ENOSPC");
expect(listDirectory()).toEqual([]);
});
});
|