whisper_web/tests/modelCache.test.ts

import { describe, test, expect } from "bun:test";
import {
  deleteModelFromCache,
  getModelStorageInfo,
  matchCachedModels,
  modelIdForUrl,
  type StorageCache,
  type StorageCacheStorage,
} from "../src/lib/modelCache";

// Pure URL matching (D-005 downloaded badge); the `caches` wrapper is
// browser-only and exercised manually. The M6 storage helpers take an
// injectable CacheStorage-shaped object, mocked below.

const TINY = "onnx-community/whisper-tiny.en";
const BASE = "onnx-community/whisper-base";

function hubUrl(modelId: string, file: string): string {
  return `https://huggingface.co/${modelId}/resolve/main/${file}`;
}

describe("matchCachedModels", () => {
  test("matches models with cached .onnx weights", () => {
    const urls = [
      hubUrl(TINY, "onnx/encoder_model_quantized.onnx"),
      hubUrl(TINY, "onnx/decoder_model_merged_quantized.onnx"),
      hubUrl(TINY, "tokenizer.json"),
    ];
    expect(matchCachedModels(urls, [TINY, BASE])).toEqual(new Set([TINY]));
  });

  test("configs/tokenizer alone do not count as downloaded", () => {
    const urls = [hubUrl(BASE, "tokenizer.json"), hubUrl(BASE, "config.json")];
    expect(matchCachedModels(urls, [TINY, BASE])).toEqual(new Set());
  });

  test("ignores query strings when checking the .onnx extension", () => {
    const urls = [`${hubUrl(BASE, "onnx/encoder_model_q4.onnx")}?download=true`];
    expect(matchCachedModels(urls, [BASE])).toEqual(new Set([BASE]));
  });

  test("does not confuse '.en' variants with their multilingual sibling", () => {
    // "onnx-community/whisper-tiny.en/..." contains the substring
    // "whisper-tiny" but must only match the .en id thanks to the trailing /.
    const urls = [hubUrl(TINY, "onnx/encoder_model_quantized.onnx")];
    expect(
      matchCachedModels(urls, [TINY, "onnx-community/whisper-tiny"]),
    ).toEqual(new Set([TINY]));
  });

  test("empty inputs yield an empty set", () => {
    expect(matchCachedModels([], [TINY])).toEqual(new Set());
    expect(matchCachedModels([hubUrl(TINY, "onnx/encoder_model.onnx")], [])).toEqual(
      new Set(),
    );
  });
});

describe("modelIdForUrl", () => {
  test("matches on the /modelId/ segment, any file type", () => {
    expect(modelIdForUrl(hubUrl(TINY, "tokenizer.json"), [TINY, BASE])).toBe(TINY);
    expect(modelIdForUrl(hubUrl(BASE, "onnx/encoder_model_q4.onnx"), [TINY, BASE])).toBe(
      BASE,
    );
  });

  test("'.en' variant URLs never match the multilingual sibling", () => {
    const url = hubUrl(TINY, "onnx/encoder_model.onnx");
    expect(modelIdForUrl(url, ["onnx-community/whisper-tiny"])).toBeUndefined();
  });

  test("unmatched URLs yield undefined", () => {
    expect(modelIdForUrl("https://example.com/other.bin", [TINY])).toBeUndefined();
  });
});

// ---------------------------------------------------------------------------
// M6 storage helpers, against a mocked Cache Storage
// ---------------------------------------------------------------------------

/** Mutable mock: cache name → (url → body size in bytes). */
function mockCacheStorage(
  data: Record<string, Record<string, number>>,
): StorageCacheStorage {
  const stores = new Map<string, Map<string, number>>(
    Object.entries(data).map(([name, entries]) => [
      name,
      new Map(Object.entries(entries)),
    ]),
  );
  function openCache(name: string): StorageCache {
    const entries = stores.get(name) ?? new Map<string, number>();
    return {
      keys: async () => [...entries.keys()].map((url) => ({ url })),
      match: async (request) => {
        const size = entries.get(request.url);
        if (size === undefined) return undefined;
        return { blob: async () => new Blob([new Uint8Array(size)]) };
      },
      delete: async (request) => entries.delete(request.url),
    };
  }
  return {
    keys: async () => [...stores.keys()],
    open: async (name) => openCache(name),
  };
}

describe("getModelStorageInfo", () => {
  test("sums bytes and counts files per model, in catalog order", async () => {
    const storage = mockCacheStorage({
      "transformers-cache": {
        [hubUrl(BASE, "onnx/encoder_model_q8.onnx")]: 2_000,
        [hubUrl(TINY, "onnx/encoder_model_quantized.onnx")]: 1_000,
        [hubUrl(TINY, "onnx/decoder_model_merged_quantized.onnx")]: 3_000,
        [hubUrl(TINY, "tokenizer.json")]: 50,
      },
    });
    expect(await getModelStorageInfo([TINY, BASE], storage)).toEqual([
      { modelId: TINY, fileCount: 3, bytes: 4_050 },
      { modelId: BASE, fileCount: 1, bytes: 2_000 },
    ]);
  });

  test("ignores non-transformers caches and unmatched URLs", async () => {
    const storage = mockCacheStorage({
      "whisper-runtime-v1": { "https://example.com/ort/ort-wasm.wasm": 9_999_999 },
      "transformers-cache": {
        [hubUrl(TINY, "config.json")]: 10,
        "https://example.com/unrelated.bin": 123,
      },
    });
    expect(await getModelStorageInfo([TINY, BASE], storage)).toEqual([
      { modelId: TINY, fileCount: 1, bytes: 10 },
    ]);
  });

  test("returns [] with no cache storage (non-browser contexts)", async () => {
    expect(await getModelStorageInfo([TINY], undefined)).toEqual([]);
  });
});

describe("deleteModelFromCache", () => {
  test("removes only the target model's entries and reports deletion", async () => {
    const storage = mockCacheStorage({
      "transformers-cache": {
        [hubUrl(TINY, "onnx/encoder_model_quantized.onnx")]: 1_000,
        [hubUrl(TINY, "tokenizer.json")]: 50,
        [hubUrl(BASE, "onnx/encoder_model_q8.onnx")]: 2_000,
      },
    });
    expect(await deleteModelFromCache(TINY, storage)).toBe(true);
    expect(await getModelStorageInfo([TINY, BASE], storage)).toEqual([
      { modelId: BASE, fileCount: 1, bytes: 2_000 },
    ]);
  });

  test("does not delete the '.en' sibling's files", async () => {
    const multilingualTiny = "onnx-community/whisper-tiny";
    const storage = mockCacheStorage({
      "transformers-cache": {
        [hubUrl(TINY, "onnx/encoder_model_quantized.onnx")]: 1_000,
      },
    });
    expect(await deleteModelFromCache(multilingualTiny, storage)).toBe(false);
    expect(await getModelStorageInfo([TINY], storage)).toEqual([
      { modelId: TINY, fileCount: 1, bytes: 1_000 },
    ]);
  });

  test("returns false when nothing matched or storage is unavailable", async () => {
    const storage = mockCacheStorage({ "transformers-cache": {} });
    expect(await deleteModelFromCache(TINY, storage)).toBe(false);
    expect(await deleteModelFromCache(TINY, undefined)).toBe(false);
  });
});