Alfred/tests/web/hooks/use-chat-scroll.test.ts

import type { UIMessage } from "ai";

import { act, renderHook } from "@testing-library/react";
import { describe, expect, mock, test } from "bun:test";

import { useChatScroll } from "../../../src/web/consoles/workbench/components/chat/use-chat-scroll";

interface HookProps {
  loadingHistory: boolean;
  messages: UIMessage[];
  scrollRef: React.RefObject<HTMLDivElement | null>;
  status: string;
  viewportElement: HTMLDivElement | null;
}

interface MockedElement {
  addEventListener: ReturnType<typeof mock>;
  clientHeight: number;
  removeEventListener: ReturnType<typeof mock>;
  scrollHeight: number;
  scrollTo: ReturnType<typeof mock>;
  scrollTop: number;
}

function asRef(el: MockedElement): React.RefObject<HTMLDivElement | null> {
  return { current: el as unknown as HTMLDivElement };
}

function asViewport(el: MockedElement): HTMLDivElement {
  return el as unknown as HTMLDivElement;
}

function createScrollElement(overrides: Partial<MockedElement> = {}): MockedElement {
  return {
    addEventListener: mock(() => {}),
    clientHeight: 400,
    removeEventListener: mock(() => {}),
    scrollHeight: 1000,
    scrollTo: mock(() => {}),
    scrollTop: 600,
    ...overrides,
  };
}

describe("useChatScroll", () => {
  test("在底部附近时自动滚动", () => {
    const el = createScrollElement({ scrollHeight: 1000, scrollTop: 920 });
    const scrollRef = asRef(el);
    const viewportElement = asViewport(el);
    const messages: UIMessage[] = [{ id: "1", parts: [], role: "user" }] as UIMessage[];

    const { rerender } = renderHook(
      ({ loadingHistory, messages, scrollRef, status, viewportElement }: HookProps) =>
        useChatScroll({ loadingHistory, messages, scrollRef, status, viewportElement }),
      {
        initialProps: { loadingHistory: false, messages, scrollRef, status: "streaming", viewportElement },
      },
    );

    const newMessages: UIMessage[] = [
      { id: "1", parts: [], role: "user" },
      { id: "2", parts: [], role: "assistant" },
    ] as UIMessage[];

    rerender({ loadingHistory: false, messages: newMessages, scrollRef, status: "streaming", viewportElement });

    expect(el.scrollTo).toHaveBeenCalled();
  });

  test("不在底部时不自动滚动", () => {
    const el = createScrollElement({ scrollHeight: 1000, scrollTop: 200 });
    const scrollRef = asRef(el);
    const viewportElement = asViewport(el);
    const messages: UIMessage[] = [{ id: "1", parts: [], role: "user" }] as UIMessage[];

    const { rerender } = renderHook(
      ({ loadingHistory, messages, scrollRef, status, viewportElement }: HookProps) =>
        useChatScroll({ loadingHistory, messages, scrollRef, status, viewportElement }),
      {
        initialProps: { loadingHistory: false, messages, scrollRef, status: "streaming", viewportElement },
      },
    );

    const handlerCalls = el.addEventListener.mock.calls as unknown[][];
    const scrollHandler = handlerCalls.find((c) => c[0] === "scroll")?.[1];

    if (scrollHandler) {
      el.scrollTop = 200;
      (scrollHandler as () => void)();
    }

    el.scrollTo.mockClear();

    const newMessages: UIMessage[] = [
      { id: "1", parts: [], role: "user" },
      { id: "2", parts: [], role: "assistant" },
    ] as UIMessage[];

    rerender({ loadingHistory: false, messages: newMessages, scrollRef, status: "streaming", viewportElement });

    const scrollToCalls = el.scrollTo.mock.calls as unknown[][];
    expect(scrollToCalls.length).toBe(0);
  });

  test("scrollToBottom 使用 instant 滚动", () => {
    const el = createScrollElement({ scrollHeight: 1000, scrollTop: 200 });
    const scrollRef = asRef(el);
    const viewportElement = asViewport(el);
    const messages: UIMessage[] = [];

    const { result } = renderHook(() => useChatScroll({ messages, scrollRef, status: "ready", viewportElement }));

    act(() => {
      result.current.scrollToBottom();
    });

    expect(el.scrollTo).toHaveBeenCalledWith({ behavior: "instant", top: 1000 });
  });

  test("流式时使用 instant 滚动", () => {
    const el = createScrollElement({ scrollHeight: 1000, scrollTop: 920 });
    const scrollRef = asRef(el);
    const viewportElement = asViewport(el);
    const messages: UIMessage[] = [{ id: "1", parts: [], role: "user" }] as UIMessage[];

    renderHook(() =>
      useChatScroll({ loadingHistory: false, messages, scrollRef, status: "streaming", viewportElement }),
    );

    const calls = el.scrollTo.mock.calls as unknown[][];
    const instantCall = calls.find((call) => {
      const opts = call[0] as Record<string, unknown> | undefined;
      return opts?.["behavior"] === "instant";
    });
    expect(instantCall).toBeTruthy();
  });

  test("loadingHistory 从 true 变为 false 时强制 scrollToBottom", async () => {
    const el = createScrollElement({ scrollHeight: 1000, scrollTop: 0 });
    const scrollRef = asRef(el);
    const viewportElement = asViewport(el);
    const messages: UIMessage[] = [{ id: "1", parts: [], role: "user" }] as UIMessage[];

    const { rerender } = renderHook(
      ({ loadingHistory, messages, scrollRef, status, viewportElement }: HookProps) =>
        useChatScroll({ loadingHistory, messages, scrollRef, status, viewportElement }),
      {
        initialProps: { loadingHistory: true, messages, scrollRef, status: "ready", viewportElement },
      },
    );

    el.scrollTo.mockClear();

    act(() => {
      rerender({ loadingHistory: false, messages, scrollRef, status: "ready", viewportElement });
    });

    await act(async () => {
      await new Promise((resolve) => setTimeout(resolve, 50));
    });

    expect(el.scrollTo).toHaveBeenCalledWith({ behavior: "instant", top: el.scrollHeight });
  });

  test("loadingHistory 不变时不触发强制滚动", () => {
    const el = createScrollElement({ scrollHeight: 1000, scrollTop: 920 });
    const scrollRef = asRef(el);
    const viewportElement = asViewport(el);
    const messages: UIMessage[] = [];

    const { rerender } = renderHook(
      ({ loadingHistory, messages, scrollRef, status, viewportElement }: HookProps) =>
        useChatScroll({ loadingHistory, messages, scrollRef, status, viewportElement }),
      {
        initialProps: { loadingHistory: false, messages, scrollRef, status: "ready", viewportElement },
      },
    );

    el.scrollTo.mockClear();

    const newMessages: UIMessage[] = [{ id: "1", parts: [], role: "user" }] as UIMessage[];
    rerender({ loadingHistory: false, messages: newMessages, scrollRef, status: "ready", viewportElement });

    const instantCalls = el.scrollTo.mock.calls.filter((call: unknown[]) => {
      const opts = call[0] as Record<string, unknown> | undefined;
      return opts?.["behavior"] === "instant" && opts?.["top"] === el.scrollHeight;
    });
    expect(instantCalls.length).toBe(1);
  });
});