ForgeLM/tauri-app/src/features/runtime/useModelDownload.test.tsx

/* @vitest-environment jsdom */

import { createRoot, type Root } from "react-dom/client";
import { act } from "react-dom/test-utils";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { HUGGING_FACE_DOWNLOAD_EVENT, type HuggingFaceDownloadProgress } from "./api";
import { useModelDownload } from "./useModelDownload";

const invokeMock = vi.fn();
const listenMock = vi.fn();

vi.mock("@tauri-apps/api/core", () => ({
  invoke: (...args: unknown[]) => invokeMock(...args),
}));

vi.mock("@tauri-apps/api/event", () => ({
  listen: (...args: unknown[]) => listenMock(...args),
}));

function Harness({ token = "hf_token" }: { token?: string }) {
  const state = useModelDownload(token);

  return (
    <div>
      <output id="is-downloading">{String(state.isDownloading)}</output>
      <output id="progress">{state.progress.toFixed(1)}</output>
      <output id="speed">{state.speed.toFixed(1)}</output>
      <output id="downloaded">{state.downloaded.toFixed(1)}</output>
      <output id="total">{state.total.toFixed(1)}</output>
      <output id="active-file">{state.activeFile ?? ""}</output>
      <output id="error">{state.error ?? ""}</output>
      <button id="start" onClick={() => void state.startDownload("microsoft/phi-4", "phi-4-q4.gguf")}>
        Start
      </button>
      <button id="cancel" onClick={() => void state.cancelDownload()}>
        Cancel
      </button>
    </div>
  );
}

describe("useModelDownload", () => {
  let container: HTMLDivElement;
  let root: Root;
  let progressListener: ((event: { payload: HuggingFaceDownloadProgress }) => void) | null = null;

  beforeEach(() => {
    // @ts-expect-error React test environment flag
    globalThis.IS_REACT_ACT_ENVIRONMENT = true;
    container = document.createElement("div");
    document.body.appendChild(container);
    root = createRoot(container);
    invokeMock.mockReset();
    listenMock.mockReset();
    progressListener = null;
    listenMock.mockImplementation(async (eventName, callback) => {
      if (eventName === HUGGING_FACE_DOWNLOAD_EVENT) {
        progressListener = callback as (event: { payload: HuggingFaceDownloadProgress }) => void;
      }
      return () => undefined;
    });
  });

  afterEach(() => {
    act(() => root.unmount());
    container.remove();
    // @ts-expect-error React test environment flag
    globalThis.IS_REACT_ACT_ENVIRONMENT = false;
  });

  it("starts a Hugging Face download with the active token and tracks progress events", async () => {
    invokeMock.mockResolvedValue("download-1");

    await act(async () => {
      root.render(<Harness token="hf_secret" />);
    });

    await act(async () => {
      (container.querySelector("#start") as HTMLButtonElement).click();
      await Promise.resolve();
    });

    expect(invokeMock).toHaveBeenCalledWith("start_hugging_face_download_command", {
      repo: "microsoft/phi-4",
      fileName: "phi-4-q4.gguf",
      token: "hf_secret",
    });
    expect(container.querySelector("#is-downloading")?.textContent).toBe("true");
    expect(container.querySelector("#active-file")?.textContent).toBe("microsoft/phi-4:phi-4-q4.gguf");

    await act(async () => {
      progressListener?.({
        payload: {
          downloadId: "download-1",
          repo: "microsoft/phi-4",
          fileName: "phi-4-q4.gguf",
          phase: "progress",
          progressPercentage: 37.5,
          downloadedMB: 768,
          totalMB: 2048,
          downloadSpeed: 12.4,
          destinationPath: null,
          message: null,
          error: null,
        },
      });
    });

    expect(container.querySelector("#progress")?.textContent).toBe("37.5");
    expect(container.querySelector("#speed")?.textContent).toBe("12.4");
    expect(container.querySelector("#downloaded")?.textContent).toBe("768.0");
    expect(container.querySelector("#total")?.textContent).toBe("2048.0");
  });

  it("ignores progress for stale downloads and cancels the active one", async () => {
    invokeMock.mockResolvedValueOnce("download-1").mockResolvedValueOnce(undefined);

    await act(async () => {
      root.render(<Harness />);
    });

    await act(async () => {
      (container.querySelector("#start") as HTMLButtonElement).click();
      await Promise.resolve();
    });

    await act(async () => {
      progressListener?.({
        payload: {
          downloadId: "download-stale",
          repo: "other/repo",
          fileName: "other.gguf",
          phase: "progress",
          progressPercentage: 90,
          downloadedMB: 90,
          totalMB: 100,
          downloadSpeed: 4,
          destinationPath: null,
          message: null,
          error: null,
        },
      });
    });

    expect(container.querySelector("#progress")?.textContent).toBe("0.0");

    await act(async () => {
      (container.querySelector("#cancel") as HTMLButtonElement).click();
      await Promise.resolve();
    });

    expect(invokeMock).toHaveBeenLastCalledWith("cancel_hugging_face_download_command", {
      downloadId: "download-1",
    });
  });
});