popiart-node-canvas/src/components/__tests__/LLMGenerateNode.test.tsx

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { render, screen, fireEvent, waitFor } from "@testing-library/react";
import { LLMGenerateNode } from "@/components/nodes/LLMGenerateNode";
import { ReactFlowProvider } from "@xyflow/react";
import { LLMGenerateNodeData } from "@/types";
import type { ProviderModel } from "@/lib/providers/types";

// Mock the workflow store
const mockUpdateNodeData = vi.fn();
const mockRegenerateNode = vi.fn();
const mockUseWorkflowStore = vi.fn();
let mockPopiLlmModels: ProviderModel[] = [];

vi.mock("@/store/workflowStore", () => ({
  useWorkflowStore: (selector?: (state: unknown) => unknown) => {
    if (selector) {
      return mockUseWorkflowStore(selector);
    }
    return mockUseWorkflowStore((s: unknown) => s);
  },
}));

vi.mock("@/store/modelStore", () => ({
  useModelStore: (selector: (state: unknown) => unknown) =>
    selector({
      byKind: {
        image: [],
        video: [],
        audio: [],
        "3d": [],
        llm: mockPopiLlmModels,
        multiAngle: [],
        highDefinition: [],
        outpainting: [],
        inpainting: [],
        videoEnhance: [],
        videoSmartErase: [],
        videoBoxErase: [],
      },
      detailsById: {},
      detailLoadingById: {},
      detailErrorById: {},
      fetchModelDetail: vi.fn(),
    }),
}));

// Wrapper component for React Flow context
function TestWrapper({ children }: { children: React.ReactNode }) {
  return <ReactFlowProvider>{children}</ReactFlowProvider>;
}

describe("LLMGenerateNode", () => {
  const originalProviderMode = process.env.NEXT_PUBLIC_PROVIDER_MODE;

  beforeEach(() => {
    vi.clearAllMocks();
    process.env.NEXT_PUBLIC_PROVIDER_MODE = "multi";
    mockPopiLlmModels = [];

    // Default mock implementation
    mockUseWorkflowStore.mockImplementation((selector) => {
      const state = {
        updateNodeData: mockUpdateNodeData,
        regenerateNode: mockRegenerateNode,
        isRunning: false,
        runningNodeIds: new Set<string>(),
        currentNodeIds: [],
        groups: {},
        nodes: [],
        getNodesWithComments: vi.fn(() => []),
        markCommentViewed: vi.fn(),
        setNavigationTarget: vi.fn(),
      };
      return selector(state);
    });
  });

  afterEach(() => {
    if (originalProviderMode === undefined) {
      delete process.env.NEXT_PUBLIC_PROVIDER_MODE;
    } else {
      process.env.NEXT_PUBLIC_PROVIDER_MODE = originalProviderMode;
    }
    window.localStorage.removeItem("node-banana-inline-parameters");
    vi.restoreAllMocks();
  });

  const createNodeData = (overrides: Partial<LLMGenerateNodeData> = {}): LLMGenerateNodeData => ({
    inputPrompt: null,
    inputImages: [],
    inputVideos: [],
    outputText: null,
    provider: "newapiwg",
    model: "doubao-seed-2-0-lite-260428",
    temperature: 1.0,
    maxTokens: 2048,
    status: "idle",
    error: null,
    ...overrides,
  });

  const createNodeProps = (data: Partial<LLMGenerateNodeData> = {}) => ({
    id: "test-llm-1",
    type: "llmGenerate" as const,
    data: createNodeData(data),
    selected: false,
  });

  describe("Basic Rendering", () => {
    it("should render generic input handle on left", () => {
      const { container } = render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps()} />
        </TestWrapper>
      );

      const inputHandle = container.querySelector('[data-handletype="generic-input"][class*="target"]');
      expect(inputHandle).toBeInTheDocument();
    });

    it("should render text output handle on right", () => {
      const { container } = render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps()} />
        </TestWrapper>
      );

      const outputHandle = container.querySelector('[data-handletype="text"][class*="source"]');
      expect(outputHandle).toBeInTheDocument();
    });

    it("hides model selectors and inline config in the node", () => {
      process.env.NEXT_PUBLIC_PROVIDER_MODE = "popi";
      window.localStorage.setItem("node-banana-inline-parameters", "true");
      mockPopiLlmModels = [{
        id: "32",
        name: "speech-llm",
        description: null,
        provider: "popiserver",
        capabilities: [],
        metadata: { type: 5, subTypes: [501] },
      }];

      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps()} />
        </TestWrapper>
      );

      expect(screen.queryByText("Provider")).not.toBeInTheDocument();
      expect(screen.queryByText("Model")).not.toBeInTheDocument();
      expect(screen.queryByText(/Temperature:/)).not.toBeInTheDocument();
      expect(screen.queryByText(/Max Tokens:/)).not.toBeInTheDocument();
      expect(screen.queryByText("温度")).not.toBeInTheDocument();
      expect(screen.queryByText("最大 Token")).not.toBeInTheDocument();
    });

    it("syncs Popi LLM model ids from modelStore without rendering a model dropdown", async () => {
      process.env.NEXT_PUBLIC_PROVIDER_MODE = "popi";
      window.localStorage.setItem("node-banana-inline-parameters", "true");
      mockPopiLlmModels = [{
        id: "37",
        name: "Kimi Server Model",
        description: null,
        provider: "popiserver",
        capabilities: [],
        metadata: { type: 5, subTypes: [501] },
      }];

      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps()} />
        </TestWrapper>
      );

      expect(screen.queryByDisplayValue("Kimi Server Model")).not.toBeInTheDocument();
      await waitFor(() => {
        expect(mockUpdateNodeData).toHaveBeenCalledWith("test-llm-1", expect.objectContaining({
          provider: "popiserver",
          model: "37",
          selectedModel: expect.objectContaining({ modelId: "37" }),
        }));
      });
    });
  });

  describe("Idle State", () => {
    it("should show 'Run to generate' message when idle and no output", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ status: "idle", outputText: null })} />
        </TestWrapper>
      );

      expect(screen.getByText("Run to generate")).toBeInTheDocument();
    });
  });

  describe("Loading State", () => {
    it("should show loading spinner when status is loading", () => {
      const { container } = render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ status: "loading" })} />
        </TestWrapper>
      );

      const spinner = container.querySelector(".animate-spin");
      expect(spinner).toBeInTheDocument();
    });
  });

  describe("Error State", () => {
    it("should show error message when status is error", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ status: "error", error: "API rate limit exceeded" })} />
        </TestWrapper>
      );

      expect(screen.getByText("API rate limit exceeded")).toBeInTheDocument();
    });

    it("should show the switched model display name for provider overload errors", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({
            status: "error",
            model: "doubao-seed-2-0-pro-260215",
            error: "doubao-seed-2-0-lite-260428: provider is temporarily overloaded. Please retry in a moment or switch models.",
          })} />
        </TestWrapper>
      );

      expect(screen.getByText((content) =>
        content.includes("Doubao Seed 2.0 Pro") &&
        (content.includes("overloaded") || content.includes("过载"))
      )).toBeInTheDocument();
    });

    it("should show 'Generation failed' when error message is null", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ status: "error", error: null })} />
        </TestWrapper>
      );

      expect(screen.getByText((content) =>
        content === "Generation failed" || content === "生成失败"
      )).toBeInTheDocument();
    });
  });

  describe("Output Text Display", () => {
    it("should display output text when data.outputText exists", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ outputText: "Generated response text" })} />
        </TestWrapper>
      );

      expect(screen.getByText("Generated response text")).toBeInTheDocument();
    });

    it("should render regenerate button when output exists", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ outputText: "Some output" })} />
        </TestWrapper>
      );

      const regenerateButton = screen.getByTitle("Regenerate");
      expect(regenerateButton).toBeInTheDocument();
    });

    it("should call regenerateNode when regenerate button is clicked", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ outputText: "Some output" })} />
        </TestWrapper>
      );

      const regenerateButton = screen.getByTitle("Regenerate");
      fireEvent.click(regenerateButton);

      expect(mockRegenerateNode).toHaveBeenCalledWith("test-llm-1");
    });

    it("should disable regenerate button when workflow is running", () => {
      mockUseWorkflowStore.mockImplementation((selector) => {
        const state = {
          updateNodeData: mockUpdateNodeData,
          regenerateNode: mockRegenerateNode,
          isRunning: true,
          runningNodeIds: new Set<string>(["test-llm-1"]),
          currentNodeIds: [],
          groups: {},
          nodes: [],
          getNodesWithComments: vi.fn(() => []),
          markCommentViewed: vi.fn(),
          setNavigationTarget: vi.fn(),
        };
        return selector(state);
      });

      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ outputText: "Some output" })} />
        </TestWrapper>
      );

      const regenerateButton = screen.getByTitle("Regenerate");
      expect(regenerateButton).toBeDisabled();
    });
  });

  describe("Clear Output Button", () => {
    it("should render clear output button when output exists", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ outputText: "Some output" })} />
        </TestWrapper>
      );

      const clearButton = screen.getByTitle("Clear output");
      expect(clearButton).toBeInTheDocument();
    });

    it("should call updateNodeData to clear output when clear button is clicked", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ outputText: "Some output" })} />
        </TestWrapper>
      );

      const clearButton = screen.getByTitle("Clear output");
      fireEvent.click(clearButton);

      expect(mockUpdateNodeData).toHaveBeenCalledWith("test-llm-1", {
        outputText: null,
        status: "idle",
        error: null,
      });
    });
  });

});