huskies/frontend/src/components/Chat.test.tsx

import {
	act,
	render,
	screen,
	waitFor,
} from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { api } from "../api/client";
import type { Message } from "../types";
import { Chat } from "./Chat";

// Module-level store for the WebSocket handlers captured during connect().
type WsHandlers = {
	onToken: (content: string) => void;
	onUpdate: (history: Message[]) => void;
	onSessionId: (sessionId: string) => void;
	onError: (message: string) => void;
	onActivity: (toolName: string) => void;
	onReconciliationProgress: (
		storyId: string,
		status: string,
		message: string,
	) => void;
};
let capturedWsHandlers: WsHandlers | null = null;

vi.mock("../api/client", () => {
	const api = {
		getOllamaModels: vi.fn(),
		getAnthropicApiKeyExists: vi.fn(),
		getAnthropicModels: vi.fn(),
		getModelPreference: vi.fn(),
		setModelPreference: vi.fn(),
		cancelChat: vi.fn(),
		setAnthropicApiKey: vi.fn(),
		readFile: vi.fn(),
		listProjectFiles: vi.fn(),
		botCommand: vi.fn(),
	};
	class ChatWebSocket {
		connect(handlers: WsHandlers) {
			capturedWsHandlers = handlers;
		}
		close() {}
		sendChat() {}
		cancel() {}
	}
	return { api, ChatWebSocket };
});

const mockedApi = {
	getOllamaModels: vi.mocked(api.getOllamaModels),
	getAnthropicApiKeyExists: vi.mocked(api.getAnthropicApiKeyExists),
	getAnthropicModels: vi.mocked(api.getAnthropicModels),
	getModelPreference: vi.mocked(api.getModelPreference),
	setModelPreference: vi.mocked(api.setModelPreference),
	cancelChat: vi.mocked(api.cancelChat),
	setAnthropicApiKey: vi.mocked(api.setAnthropicApiKey),
	readFile: vi.mocked(api.readFile),
	listProjectFiles: vi.mocked(api.listProjectFiles),
	botCommand: vi.mocked(api.botCommand),
};

function setupMocks() {
	mockedApi.getOllamaModels.mockResolvedValue(["llama3.1"]);
	mockedApi.getAnthropicApiKeyExists.mockResolvedValue(true);
	mockedApi.getAnthropicModels.mockResolvedValue([]);
	mockedApi.getModelPreference.mockResolvedValue(null);
	mockedApi.setModelPreference.mockResolvedValue(true);
	mockedApi.readFile.mockResolvedValue("");
	mockedApi.listProjectFiles.mockResolvedValue([]);
	mockedApi.cancelChat.mockResolvedValue(true);
	mockedApi.setAnthropicApiKey.mockResolvedValue(true);
	mockedApi.botCommand.mockResolvedValue({ response: "Bot response" });
}

describe("Default provider selection (Story 206)", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
	});

	it("AC1: defaults to claude-code-pty when no saved model preference exists", async () => {
		mockedApi.getOllamaModels.mockResolvedValue([]);
		mockedApi.getAnthropicApiKeyExists.mockResolvedValue(false);
		mockedApi.getAnthropicModels.mockResolvedValue([]);
		mockedApi.getModelPreference.mockResolvedValue(null);
		mockedApi.setModelPreference.mockResolvedValue(true);
		mockedApi.cancelChat.mockResolvedValue(true);

		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		// With no models available, the header renders a text input with the model value
		const input = screen.getByPlaceholderText("Model");
		expect(input).toHaveValue("claude-code-pty");
	});

	it("AC2: claude-code-pty remains default even when ollama models are available", async () => {
		mockedApi.getOllamaModels.mockResolvedValue(["llama3.1", "deepseek-coder"]);
		mockedApi.getAnthropicApiKeyExists.mockResolvedValue(false);
		mockedApi.getAnthropicModels.mockResolvedValue([]);
		mockedApi.getModelPreference.mockResolvedValue(null);
		mockedApi.setModelPreference.mockResolvedValue(true);
		mockedApi.cancelChat.mockResolvedValue(true);

		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		// Wait for Ollama models to load and the select dropdown to appear
		const select = await screen.findByRole("combobox");
		expect(select).toHaveValue("claude-code-pty");
	});

	it("AC3: respects saved model preference for existing projects", async () => {
		mockedApi.getOllamaModels.mockResolvedValue(["llama3.1"]);
		mockedApi.getAnthropicApiKeyExists.mockResolvedValue(false);
		mockedApi.getAnthropicModels.mockResolvedValue([]);
		mockedApi.getModelPreference.mockResolvedValue("llama3.1");
		mockedApi.setModelPreference.mockResolvedValue(true);
		mockedApi.cancelChat.mockResolvedValue(true);

		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		// Wait for models to load and preference to be applied
		const select = await screen.findByRole("combobox");
		await waitFor(() => {
			expect(select).toHaveValue("llama3.1");
		});
	});
});

describe("Chat message rendering — unified tool call UI", () => {
	beforeEach(() => {
		capturedWsHandlers = null;
		setupMocks();
	});

	it("renders tool call badge for assistant message with tool_calls (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Read src/main.rs" },
			{
				role: "assistant",
				content: "I'll read that file.",
				tool_calls: [
					{
						id: "toolu_abc",
						type: "function",
						function: {
							name: "Read",
							arguments: '{"file_path":"src/main.rs"}',
						},
					},
				],
			},
		];

		await act(async () => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(await screen.findByText("I'll read that file.")).toBeInTheDocument();
		expect(await screen.findByText("Read(src/main.rs)")).toBeInTheDocument();
	});

	it("renders collapsible tool output for tool role messages (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Check the file" },
			{
				role: "assistant",
				content: "",
				tool_calls: [
					{
						id: "toolu_1",
						type: "function",
						function: { name: "Read", arguments: '{"file_path":"foo.rs"}' },
					},
				],
			},
			{
				role: "tool",
				content: 'fn main() { println!("hello"); }',
				tool_call_id: "toolu_1",
			},
			{ role: "assistant", content: "The file contains a main function." },
		];

		await act(async () => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(await screen.findByText(/Tool Output/)).toBeInTheDocument();
		expect(
			await screen.findByText("The file contains a main function."),
		).toBeInTheDocument();
	});

	it("renders plain assistant message without tool call badges (AC5)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Hello" },
			{ role: "assistant", content: "Hi there! How can I help?" },
		];

		await act(async () => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(
			await screen.findByText("Hi there! How can I help?"),
		).toBeInTheDocument();
		expect(screen.queryByText(/Tool Output/)).toBeNull();
	});

	it("renders multiple tool calls in a single assistant turn (AC3)", async () => {
		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => expect(capturedWsHandlers).not.toBeNull());

		const messages: Message[] = [
			{ role: "user", content: "Do some work" },
			{
				role: "assistant",
				content: "I'll do multiple things.",
				tool_calls: [
					{
						id: "id1",
						type: "function",
						function: { name: "Bash", arguments: '{"command":"cargo test"}' },
					},
					{
						id: "id2",
						type: "function",
						function: { name: "Read", arguments: '{"file_path":"Cargo.toml"}' },
					},
				],
			},
		];

		await act(async () => {
			capturedWsHandlers?.onUpdate(messages);
		});

		expect(await screen.findByText("Bash(cargo test)")).toBeInTheDocument();
		expect(await screen.findByText("Read(Cargo.toml)")).toBeInTheDocument();
	});

	it("does not fetch Anthropic models when no API key exists", async () => {
		mockedApi.getAnthropicApiKeyExists.mockResolvedValue(false);
		mockedApi.getAnthropicModels.mockClear();

		render(<Chat projectPath="/tmp/project" onCloseProject={vi.fn()} />);

		await waitFor(() => {
			expect(mockedApi.getAnthropicApiKeyExists).toHaveBeenCalled();
		});

		expect(mockedApi.getAnthropicModels).not.toHaveBeenCalled();
	});
});