src/tokenizer/cache.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

import {
  fetchTokenizerAsset,
  TOKENIZER_CACHE,
  tokenizerCacheSize,
} from "./cache";
import type { RemoteFile } from "./types";

const HELLO_SHA256 =
  "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824";
const HELLO_GIT_SHA1 = "b6fc4c620b67d95f953a5c1c1230aaab5db5a1b0";
const CHANGED_SHA256 =
  "d67e2e944994496c8d8ec76eed0cf9f09679448d584b532bebf941852a37f5ed";

function file(url: string, algorithm: "git-sha1" | "sha256", value: string): RemoteFile {
  return {
    filename: "tokenizer.json",
    url,
    hash: { algorithm, value },
  };
}

describe("tokenizer asset cache", () => {
  const stores = new Map<string, Map<string, Response>>();

  beforeEach(() => {
    stores.clear();
    vi.stubGlobal("self", { location: { origin: "https://tokenize-it.test" } });
    vi.stubGlobal("caches", {
      async open(name: string) {
        let store = stores.get(name);
        if (!store) {
          store = new Map();
          stores.set(name, store);
        }
        return {
          async match(key: RequestInfo | URL) {
            const url =
              typeof key === "string"
                ? key
                : key instanceof URL
                  ? key.href
                  : key.url;
            return store.get(url)?.clone();
          },
          async put(key: string, response: Response) {
            store.set(key, response.clone());
          },
          async keys() {
            return Array.from(store.keys(), (key) => new Request(key));
          },
        };
      },
      async delete(name: string) {
        return stores.delete(name);
      },
      async keys() {
        return Array.from(stores.keys());
      },
    });
  });

  afterEach(() => {
    vi.unstubAllGlobals();
  });

  it("reuses identical content across different repository URLs", async () => {
    const fetcher = vi.fn<typeof fetch>().mockResolvedValue(
      new Response("hello", {
        headers: { "Content-Type": "application/json" },
      }),
    );
    vi.stubGlobal("fetch", fetcher);

    const first = await fetchTokenizerAsset(
      file("https://huggingface.co/owner/model-a/resolve/commit-a/tokenizer.json", "sha256", HELLO_SHA256),
      vi.fn(),
    );
    const second = await fetchTokenizerAsset(
      file("https://huggingface.co/owner/model-b/resolve/commit-b/tokenizer.json", "sha256", HELLO_SHA256),
      vi.fn(),
    );

    expect(new TextDecoder().decode(first)).toBe("hello");
    expect(new TextDecoder().decode(second)).toBe("hello");
    expect(fetcher).toHaveBeenCalledTimes(1);
    expect(stores.get(TOKENIZER_CACHE)?.size).toBe(1);
  });

  it("downloads a new record when the file hash changes", async () => {
    const fetcher = vi
      .fn<typeof fetch>()
      .mockResolvedValueOnce(new Response("hello"))
      .mockResolvedValueOnce(new Response("changed"));
    vi.stubGlobal("fetch", fetcher);

    await fetchTokenizerAsset(
      file("https://huggingface.co/owner/model/resolve/commit-a/tokenizer.json", "sha256", HELLO_SHA256),
      vi.fn(),
    );
    await fetchTokenizerAsset(
      file("https://huggingface.co/owner/model/resolve/commit-b/tokenizer.json", "sha256", CHANGED_SHA256),
      vi.fn(),
    );

    expect(fetcher).toHaveBeenCalledTimes(2);
    expect(stores.get(TOKENIZER_CACHE)?.size).toBe(2);
  });

  it("reports the total size of cached tokenizer assets", async () => {
    const fetcher = vi
      .fn<typeof fetch>()
      .mockResolvedValueOnce(new Response("hello"))
      .mockResolvedValueOnce(new Response("changed"));
    vi.stubGlobal("fetch", fetcher);

    await fetchTokenizerAsset(
      file("https://huggingface.co/owner/model/resolve/commit-a/tokenizer.json", "sha256", HELLO_SHA256),
      vi.fn(),
    );
    await fetchTokenizerAsset(
      file("https://huggingface.co/owner/model/resolve/commit-b/tokenizer.json", "sha256", CHANGED_SHA256),
      vi.fn(),
    );

    await expect(tokenizerCacheSize()).resolves.toBe(12);
  });

  it("validates regular files using the Git blob hash", async () => {
    vi.stubGlobal(
      "fetch",
      vi.fn<typeof fetch>().mockResolvedValue(new Response("hello")),
    );

    await expect(
      fetchTokenizerAsset(
        file("https://huggingface.co/owner/model/resolve/commit/tokenizer.json", "git-sha1", HELLO_GIT_SHA1),
        vi.fn(),
      ),
    ).resolves.toEqual(new TextEncoder().encode("hello"));
  });

  it("rejects downloaded content that does not match the Hub hash", async () => {
    vi.stubGlobal(
      "fetch",
      vi.fn<typeof fetch>().mockResolvedValue(new Response("changed")),
    );

    await expect(
      fetchTokenizerAsset(
        file("https://huggingface.co/owner/model/resolve/commit/tokenizer.json", "sha256", HELLO_SHA256),
        vi.fn(),
      ),
    ).rejects.toThrow("文件与预期哈希不一致");
    expect(stores.get(TOKENIZER_CACHE)?.size).toBe(0);
  });
});