frontend/src/components/__tests__/NiiVueViewer.test.tsx

import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, screen, waitFor } from "@testing-library/react";
import { NiiVueViewer } from "../NiiVueViewer";

// Store mock function references so tests can verify calls
const mockLoadVolumes = vi.fn().mockResolvedValue(undefined);
const mockCleanup = vi.fn();
const mockAttachToCanvas = vi.fn();
const mockLoseContext = vi.fn();

// Mock the NiiVue module since it requires actual WebGL
vi.mock("@niivue/niivue", () => ({
  Niivue: class MockNiivue {
    attachToCanvas = mockAttachToCanvas;
    loadVolumes = mockLoadVolumes;
    setSliceType = vi.fn();
    cleanup = mockCleanup;
    gl = {
      getExtension: vi.fn(() => ({ loseContext: mockLoseContext })),
    };
    opts = {};
  },
}));

describe("NiiVueViewer", () => {
  const defaultProps = {
    backgroundUrl: "http://localhost:7860/files/dwi.nii.gz",
  };

  beforeEach(() => {
    vi.clearAllMocks();
  });

  it("renders canvas element", () => {
    render(<NiiVueViewer {...defaultProps} />);

    expect(document.querySelector("canvas")).toBeInTheDocument();
  });

  it("renders container with correct styling", () => {
    render(<NiiVueViewer {...defaultProps} />);

    const container = document.querySelector("canvas")?.parentElement;
    expect(container).toHaveClass("bg-gray-900");
  });

  it("renders help text for controls", () => {
    render(<NiiVueViewer {...defaultProps} />);

    expect(screen.getByText(/scroll/i)).toBeInTheDocument();
    expect(screen.getByText(/drag/i)).toBeInTheDocument();
  });

  it("attaches NiiVue to canvas on mount", () => {
    render(<NiiVueViewer {...defaultProps} />);

    expect(mockAttachToCanvas).toHaveBeenCalled();
    // Verify it was called with a canvas element
    const arg = mockAttachToCanvas.mock.calls[0][0];
    expect(arg).toBeInstanceOf(HTMLCanvasElement);
  });

  it("loads background volume on mount", () => {
    render(<NiiVueViewer {...defaultProps} />);

    expect(mockLoadVolumes).toHaveBeenCalledWith([
      { url: defaultProps.backgroundUrl, colormap: "gray", opacity: 1 },
    ]);
  });

  it("loads both background and overlay when overlayUrl provided", () => {
    const overlayUrl = "http://localhost:7860/files/prediction.nii.gz";

    render(<NiiVueViewer {...defaultProps} overlayUrl={overlayUrl} />);

    expect(mockLoadVolumes).toHaveBeenCalledWith([
      { url: defaultProps.backgroundUrl, colormap: "gray", opacity: 1 },
      { url: overlayUrl, colormap: "red", opacity: 0.5 },
    ]);
  });

  it("calls cleanup on unmount", () => {
    const { unmount } = render(<NiiVueViewer {...defaultProps} />);

    unmount();

    expect(mockCleanup).toHaveBeenCalled();
    expect(mockLoseContext).toHaveBeenCalled();
  });

  it("sets canvas dimensions", () => {
    render(<NiiVueViewer {...defaultProps} />);

    const canvas = document.querySelector("canvas");
    expect(canvas).toHaveClass("w-full", "h-[500px]");
  });

  it("displays error when volume loading fails", async () => {
    const errorMessage = "Network error loading volume";
    mockLoadVolumes.mockRejectedValueOnce(new Error(errorMessage));

    render(<NiiVueViewer {...defaultProps} />);

    // Wait for error to be displayed
    const errorElement = await screen.findByText(/failed to load volume/i);
    expect(errorElement).toBeInTheDocument();
    expect(errorElement).toHaveTextContent(errorMessage);
  });

  it("calls onError callback when volume loading fails", async () => {
    const errorMessage = "Network error";
    const onError = vi.fn();
    mockLoadVolumes.mockRejectedValueOnce(new Error(errorMessage));

    render(<NiiVueViewer {...defaultProps} onError={onError} />);

    // Wait for error callback to be invoked (use RTL's waitFor, not vi.waitFor)
    await waitFor(() => {
      expect(onError).toHaveBeenCalledWith(errorMessage);
    });
  });

  it("ignores errors from stale loads after URL change", async () => {
    const onError = vi.fn();
    // First load succeeds, second load fails slowly
    let rejectSecondLoad: (error: Error) => void;
    mockLoadVolumes.mockResolvedValueOnce(undefined).mockImplementationOnce(
      () =>
        new Promise((_, reject) => {
          rejectSecondLoad = reject;
        }),
    );

    const { rerender } = render(
      <NiiVueViewer
        backgroundUrl="http://localhost/first.nii.gz"
        onError={onError}
      />,
    );

    // Change URL - starts second load
    rerender(
      <NiiVueViewer
        backgroundUrl="http://localhost/second.nii.gz"
        onError={onError}
      />,
    );

    // Change URL again - makes second load stale
    rerender(
      <NiiVueViewer
        backgroundUrl="http://localhost/third.nii.gz"
        onError={onError}
      />,
    );

    // Now reject the second load (stale)
    rejectSecondLoad!(new Error("Stale load error"));

    // Flush async work (let rejection be processed) before asserting
    // Using waitFor with negative assertions is flaky - it passes immediately
    await new Promise((resolve) => setTimeout(resolve, 0));
    expect(onError).not.toHaveBeenCalledWith("Stale load error");
  });
});