// @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); }); });