nexus/ui/src/hooks/useStreamingChat.test.ts

// @vitest-environment jsdom
import { describe, it, expect, vi, beforeEach } from "vitest";
import { renderHook, act } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { createElement } from "react";
import { useStreamingChat } from "./useStreamingChat";
import { chatApi } from "../api/chat";

// Mock chatApi
vi.mock("../api/chat", () => ({
  chatApi: {
    postMessageAndStream: vi.fn(),
    savePartialMessage: vi.fn().mockResolvedValue({}),
    postMessage: vi.fn().mockResolvedValue({}),
  },
}));

function createWrapper() {
  const queryClient = new QueryClient({
    defaultOptions: { queries: { retry: false } },
  });
  return ({ children }: { children: React.ReactNode }) =>
    createElement(QueryClientProvider, { client: queryClient }, children);
}

describe("useStreamingChat", () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  it("accumulates tokens from onToken callbacks into streamingContent", async () => {
    // Mock postMessageAndStream to capture callbacks and simulate tokens
    let capturedCallbacks: {
      onToken: (token: string) => void;
      onDone: (messageId: string, content: string) => void;
      onError: (error: string) => void;
    };
    vi.mocked(chatApi.postMessageAndStream).mockImplementation(
      (_convId, _data, callbacks, _signal) => {
        capturedCallbacks = callbacks;
        return Promise.resolve();
      },
    );

    const { result } = renderHook(
      () => useStreamingChat("conv-1"),
      { wrapper: createWrapper() },
    );

    // Start stream
    act(() => {
      result.current.startStream("hello world");
    });

    expect(result.current.isStreaming).toBe(true);
    expect(result.current.streamingContent).toBe("");

    // Simulate tokens arriving
    act(() => {
      capturedCallbacks!.onToken("Hello ");
    });
    expect(result.current.streamingContent).toBe("Hello ");

    act(() => {
      capturedCallbacks!.onToken("world!");
    });
    expect(result.current.streamingContent).toBe("Hello world!");
  });

  it("sets isStreaming=true when stream starts, false on done", async () => {
    let capturedCallbacks: {
      onToken: (token: string) => void;
      onDone: (messageId: string, content: string) => void;
      onError: (error: string) => void;
    };
    vi.mocked(chatApi.postMessageAndStream).mockImplementation(
      (_convId, _data, callbacks, _signal) => {
        capturedCallbacks = callbacks;
        return Promise.resolve();
      },
    );

    const { result } = renderHook(
      () => useStreamingChat("conv-1"),
      { wrapper: createWrapper() },
    );

    expect(result.current.isStreaming).toBe(false);

    act(() => {
      result.current.startStream("test");
    });
    expect(result.current.isStreaming).toBe(true);

    act(() => {
      capturedCallbacks!.onDone("msg-1", "test response");
    });
    expect(result.current.isStreaming).toBe(false);
    expect(result.current.streamingContent).toBe("");
  });

  it("stop() aborts the controller and sets isStreaming=false", async () => {
    let capturedSignal: AbortSignal | undefined;
    vi.mocked(chatApi.postMessageAndStream).mockImplementation(
      (_convId, _data, _callbacks, signal) => {
        capturedSignal = signal;
        return Promise.resolve();
      },
    );

    const { result } = renderHook(
      () => useStreamingChat("conv-1"),
      { wrapper: createWrapper() },
    );

    act(() => {
      result.current.startStream("test");
    });
    expect(result.current.isStreaming).toBe(true);
    expect(capturedSignal?.aborted).toBe(false);

    act(() => {
      result.current.stop();
    });
    expect(result.current.isStreaming).toBe(false);
    expect(capturedSignal?.aborted).toBe(true);
  });

  it("handles SSE error by setting isStreaming=false", async () => {
    let capturedCallbacks: {
      onToken: (token: string) => void;
      onDone: (messageId: string, content: string) => void;
      onError: (error: string) => void;
    };
    vi.mocked(chatApi.postMessageAndStream).mockImplementation(
      (_convId, _data, callbacks, _signal) => {
        capturedCallbacks = callbacks;
        return Promise.resolve();
      },
    );

    const { result } = renderHook(
      () => useStreamingChat("conv-1"),
      { wrapper: createWrapper() },
    );

    act(() => {
      result.current.startStream("test");
    });
    expect(result.current.isStreaming).toBe(true);

    act(() => {
      capturedCallbacks!.onError("Stream error");
    });
    expect(result.current.isStreaming).toBe(false);
  });

  it("does nothing when conversationId is null", () => {
    const { result } = renderHook(
      () => useStreamingChat(null),
      { wrapper: createWrapper() },
    );

    act(() => {
      result.current.startStream("test");
    });
    expect(chatApi.postMessageAndStream).not.toHaveBeenCalled();
    expect(result.current.isStreaming).toBe(false);
  });
});