feat(@huggingface/hub): adding scanCacheDir

packages/hub/package.json

@@ -20,6 +20,7 @@
 	"browser": {
 		"./src/utils/sha256-node.ts": false,
 		"./src/utils/FileBlob.ts": false,
+		"./src/lib/cache-management.ts": false,
 		"./dist/index.js": "./dist/browser/index.js",
 		"./dist/index.mjs": "./dist/browser/index.mjs"
 	},

packages/hub/src/lib/cache-management.spec.ts

@@ -0,0 +1,137 @@
+import { describe, test, expect, vi, beforeEach } from "vitest";
+import {
+	scanCacheDir,
+	scanCachedRepo,
+	REPO_TYPE_T,
+	scanSnapshotDir,
+	parseRepoType,
+	getBlobStat,
+	type CachedFileInfo,
+} from "./cache-management";
+import { stat, readdir, realpath, lstat } from "node:fs/promises";
+import type { Dirent, Stats } from "node:fs";
+import { join } from "node:path";
+
+// Mocks
+vi.mock("node:fs/promises");
+
+beforeEach(() => {
+	vi.resetAllMocks();
+	vi.restoreAllMocks();
+});
+
+describe("scanCacheDir", () => {
+	test("should throw an error if cacheDir is not a directory", async () => {
+		vi.mocked(stat).mockResolvedValueOnce({
+			isDirectory: () => false,
+		} as Stats);
+
+		await expect(scanCacheDir("/fake/dir")).rejects.toThrow("Scan cache expects a directory");
+	});
+
+	test("empty directory should return an empty set of repository and no warnings", async () => {
+		vi.mocked(stat).mockResolvedValueOnce({
+			isDirectory: () => true,
+		} as Stats);
+
+		// mock empty cache folder
+		vi.mocked(readdir).mockResolvedValue([]);
+
+		const result = await scanCacheDir("/fake/dir");
+
+		// cacheDir must have been read
+		expect(readdir).toHaveBeenCalledWith("/fake/dir");
+
+		expect(result.warnings.length).toBe(0);
+		expect(result.repos.size).toBe(0);
+		expect(result.sizeOnDisk).toBe(0);
+	});
+});
+
+describe("scanCachedRepo", () => {
+	test("should throw an error for invalid repo path", async () => {
+		await expect(() => {
+			return scanCachedRepo("/fake/repo_path");
+		}).rejects.toThrow("Repo path is not a valid HuggingFace cache directory");
+	});
+
+	test("should throw an error if the snapshot folder does not exist", async () => {
+		vi.mocked(readdir).mockResolvedValue([]);
+		vi.mocked(stat).mockResolvedValue({
+			isDirectory: () => false,
+		} as Stats);
+
+		await expect(() => {
+			return scanCachedRepo("/fake/cacheDir/models--hello-world--name");
+		}).rejects.toThrow("Snapshots dir doesn't exist in cached repo");
+	});
+
+	test("should properly parse the repository name", async () => {
+		const repoPath = "/fake/cacheDir/models--hello-world--name";
+		vi.mocked(readdir).mockResolvedValue([]);
+		vi.mocked(stat).mockResolvedValue({
+			isDirectory: () => true,
+		} as Stats);
+
+		const result = await scanCachedRepo(repoPath);
+		expect(readdir).toHaveBeenCalledWith(join(repoPath, "refs"), {
+			withFileTypes: true,
+		});
+
+		expect(result.repoId).toBe("hello-world/name");
+	});
+});
+
+describe("scanSnapshotDir", () => {
+	test("should scan a valid snapshot directory", async () => {
+		const cachedFiles = new Set<CachedFileInfo>();
+		const blobStats = new Map<string, Stats>();
+		vi.mocked(readdir).mockResolvedValueOnce([{ name: "file1", isDirectory: () => false } as Dirent]);
+
+		vi.mocked(realpath).mockResolvedValueOnce("/fake/realpath");
+		vi.mocked(lstat).mockResolvedValueOnce({ size: 1024, atimeMs: Date.now(), mtimeMs: Date.now() } as Stats);
+
+		await scanSnapshotDir("/fake/revision", cachedFiles, blobStats);
+
+		expect(cachedFiles.size).toBe(1);
+		expect(blobStats.size).toBe(1);
+	});
+});
+
+describe("getBlobStat", () => {
+	test("should retrieve blob stat if already cached", async () => {
+		const blobStats = new Map<string, Stats>([["/fake/blob", { size: 1024 } as Stats]]);
+		const result = await getBlobStat("/fake/blob", blobStats);
+
+		expect(lstat).not.toHaveBeenCalled();
+		expect(result.size).toBe(1024);
+	});
+
+	test("should fetch and cache blob stat if not cached", async () => {
+		const blobStats = new Map();
+		vi.mocked(lstat).mockResolvedValueOnce({ size: 2048 } as Stats);
+
+		const result = await getBlobStat("/fake/blob", blobStats);
+
+		expect(result.size).toBe(2048);
+		expect(blobStats.size).toBe(1);
+	});
+});
+
+describe("parseRepoType", () => {
+	test("should parse models repo type", () => {
+		expect(parseRepoType("models")).toBe(REPO_TYPE_T.MODEL);
+	});
+
+	test("should parse model repo type", () => {
+		expect(parseRepoType("model")).toBe(REPO_TYPE_T.MODEL);
+	});
+
+	test("should parse dataset repo type", () => {
+		expect(parseRepoType("dataset")).toBe(REPO_TYPE_T.DATASET);
+	});
+
+	test("should throw an error for invalid repo type", () => {
+		expect(() => parseRepoType("invalid")).toThrow();
+	});
+});