whisper_web/tests/protocol.test.ts

import { describe, test, expect } from "bun:test";
import {
  applyLoadProgress,
  initialLoadProgress,
  isUiToWorkerMessage,
  isWorkerToUiMessage,
  type LoadProgressMessage,
  type UiToWorkerMessage,
  type WorkerToUiMessage,
} from "../src/worker/protocol";

// Protocol handling only — no worker, no network, no model weights.

describe("isUiToWorkerMessage", () => {
  test("accepts valid load and transcribe messages", () => {
    const loads: UiToWorkerMessage[] = [
      { type: "load", modelId: "onnx-community/whisper-tiny.en", device: "wasm" },
      { type: "load", modelId: "onnx-community/whisper-tiny", device: "webgpu" },
      { type: "load", modelId: "onnx-community/whisper-base", device: "auto" },
      { type: "load", modelId: "onnx-community/whisper-base", device: "auto", dtype: "q4" },
      {
        type: "load",
        modelId: "onnx-community/whisper-large-v3-turbo",
        device: "webgpu",
        dtype: { encoder_model: "fp16", decoder_model_merged: "q4" },
      },
    ];
    const transcribes: UiToWorkerMessage[] = [
      { type: "transcribe", audio: new Float32Array(8) },
      { type: "transcribe", audio: new Float32Array(8), language: "da", task: "transcribe" },
      { type: "transcribe", audio: new Float32Array(8), language: "auto", task: "translate" },
    ];
    for (const message of [...loads, ...transcribes]) {
      expect(isUiToWorkerMessage(message)).toBe(true);
    }
  });

  test("rejects malformed messages", () => {
    expect(isUiToWorkerMessage(null)).toBe(false);
    expect(isUiToWorkerMessage("load")).toBe(false);
    expect(isUiToWorkerMessage({ type: "load" })).toBe(false);
    expect(isUiToWorkerMessage({ type: "load", modelId: "x", device: "cuda" })).toBe(false);
    expect(isUiToWorkerMessage({ type: "load", modelId: "x", device: "auto", dtype: "int3" })).toBe(false);
    expect(
      isUiToWorkerMessage({
        type: "load",
        modelId: "x",
        device: "auto",
        dtype: { encoder_model: "fp16" },
      }),
    ).toBe(false);
    expect(isUiToWorkerMessage({ type: "transcribe", audio: [1, 2, 3] })).toBe(false);
    expect(
      isUiToWorkerMessage({ type: "transcribe", audio: new Float32Array(4), task: "detect" }),
    ).toBe(false);
    expect(isUiToWorkerMessage({ type: "result", text: "", chunks: [] })).toBe(false);
  });
});

describe("isWorkerToUiMessage", () => {
  test("accepts every worker → UI message shape", () => {
    const messages: WorkerToUiMessage[] = [
      { type: "load-progress", status: "initiate", file: "model.onnx" },
      { type: "load-progress", status: "download", file: "model.onnx" },
      {
        type: "load-progress",
        status: "progress",
        file: "model.onnx",
        loaded: 10,
        total: 100,
        progress: 10,
      },
      { type: "load-progress", status: "done", file: "model.onnx" },
      { type: "load-progress", status: "ready", device: "wasm" },
      { type: "load-progress", status: "ready", device: "webgpu" },
      { type: "partial", text: " hello" },
      { type: "result", text: "hello world", chunks: [{ text: "hello world", timestamp: [0, 1.2] }] },
      { type: "error", stage: "load", message: "boom" },
    ];
    for (const message of messages) {
      expect(isWorkerToUiMessage(message)).toBe(true);
    }
  });

  test("rejects malformed messages", () => {
    expect(isWorkerToUiMessage(undefined)).toBe(false);
    expect(isWorkerToUiMessage({})).toBe(false);
    expect(isWorkerToUiMessage({ type: "load-progress", status: "progress", file: "f" })).toBe(false);
    expect(isWorkerToUiMessage({ type: "load-progress", status: "ready" })).toBe(false);
    expect(isWorkerToUiMessage({ type: "load-progress", status: "ready", device: "cuda" })).toBe(false);
    expect(isWorkerToUiMessage({ type: "partial" })).toBe(false);
    expect(isWorkerToUiMessage({ type: "result", text: "x", chunks: "nope" })).toBe(false);
    expect(isWorkerToUiMessage({ type: "error", stage: "decode", message: "m" })).toBe(false);
    expect(isWorkerToUiMessage({ type: "load", modelId: "x", device: "wasm" })).toBe(false);
  });
});

describe("applyLoadProgress", () => {
  test("tracks per-file byte progress without mutating previous state", () => {
    const p1: LoadProgressMessage = { type: "load-progress", status: "initiate", file: "a.onnx" };
    const p2: LoadProgressMessage = {
      type: "load-progress",
      status: "progress",
      file: "a.onnx",
      loaded: 50,
      total: 200,
      progress: 25,
    };

    const s1 = applyLoadProgress(initialLoadProgress, p1);
    const s2 = applyLoadProgress(s1, p2);

    expect(initialLoadProgress.files).toEqual({});
    expect(s1.files["a.onnx"]).toEqual({ file: "a.onnx", loaded: 0, total: 0, done: false });
    expect(s1.files["a.onnx"]!.done).toBe(false); // s1 untouched by s2
    expect(s2.files["a.onnx"]).toEqual({ file: "a.onnx", loaded: 50, total: 200, done: false });
    expect(s2.ready).toBe(false);
  });

  test("tracks multiple files independently", () => {
    let state = initialLoadProgress;
    state = applyLoadProgress(state, {
      type: "load-progress",
      status: "progress",
      file: "a.onnx",
      loaded: 10,
      total: 100,
      progress: 10,
    });
    state = applyLoadProgress(state, {
      type: "load-progress",
      status: "progress",
      file: "b.json",
      loaded: 5,
      total: 5,
      progress: 100,
    });
    expect(Object.keys(state.files)).toEqual(["a.onnx", "b.json"]);
    expect(state.files["a.onnx"]!.loaded).toBe(10);
    expect(state.files["b.json"]!.total).toBe(5);
  });

  test("done snaps a file's loaded bytes to its total", () => {
    let state = applyLoadProgress(initialLoadProgress, {
      type: "load-progress",
      status: "progress",
      file: "a.onnx",
      loaded: 199,
      total: 200,
      progress: 99.5,
    });
    state = applyLoadProgress(state, { type: "load-progress", status: "done", file: "a.onnx" });
    expect(state.files["a.onnx"]).toEqual({ file: "a.onnx", loaded: 200, total: 200, done: true });
  });

  test("ready flips the ready flag, records the device, keeps file state", () => {
    let state = applyLoadProgress(initialLoadProgress, {
      type: "load-progress",
      status: "done",
      file: "a.onnx",
    });
    expect(state.device).toBeNull();
    state = applyLoadProgress(state, {
      type: "load-progress",
      status: "ready",
      device: "webgpu",
    });
    expect(state.ready).toBe(true);
    expect(state.device).toBe("webgpu");
    expect(state.files["a.onnx"]!.done).toBe(true);
  });
});