from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest from typer.testing import CliRunner from haiku.rag.cli import _cli as cli from haiku.rag.skills.rag import RAGState runner = CliRunner() def test_chat_command(): """Test chat command launches chat TUI.""" with patch("haiku.rag.chat.run_chat") as mock_chat: mock_chat.return_value = None result = runner.invoke(cli, ["chat"]) assert result.exit_code == 0 mock_chat.assert_called_once() def test_run_chat_creates_app_and_runs(temp_db_path: Path): """Test run_chat() creates a ChatApp and calls run().""" with patch("haiku.rag.chat.app.ChatApp.run") as mock_run: from haiku.rag.chat import run_chat run_chat(db_path=temp_db_path) mock_run.assert_called_once() def _make_mock_client(): """Create a mock HaikuRAG client.""" mock_client = AsyncMock() mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aexit__ = AsyncMock(return_value=None) return mock_client def _make_app(db_path: Path, mock_client: AsyncMock | None = None): """Create a ChatApp with mocked HaikuRAG.""" from haiku.rag.chat.app import ChatApp if mock_client is None: mock_client = _make_mock_client() skill = MagicMock() skill.state_type = None skill.state_namespace = None skill.tools = [] skill.toolsets = [] skill.resources = [] skill.metadata = MagicMock() skill.metadata.name = "rag" skill.metadata.description = "RAG skill" return ChatApp( db_path=db_path, skills=[skill], read_only=True, ), mock_client def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None): """Create a ChatApp with a skill that has RAGState.""" from haiku.rag.chat.app import ChatApp from haiku.skills.models import Skill, SkillMetadata, SkillSource if mock_client is None: mock_client = _make_mock_client() skill = Skill( metadata=SkillMetadata(name="rag", description="RAG skill"), source=SkillSource.ENTRYPOINT, tools=[], state_type=RAGState, state_namespace="rag", ) return ChatApp( db_path=db_path, skills=[skill], read_only=True, ), mock_client @pytest.mark.asyncio async def test_chat_app_has_required_widgets(temp_db_path: Path): """Test that ChatApp has the required widgets: ChatHistory, Input.""" from haiku.rag.chat.widgets.chat_history import ChatHistory app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test(): chat_history = app.query_one(ChatHistory) assert chat_history is not None from textual.widgets import Input chat_input = app.query_one(Input) assert chat_input is not None @pytest.mark.asyncio async def test_chat_app_quit_binding(temp_db_path: Path): """Test that pressing ctrl+q quits the app.""" app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test() as pilot: assert app.is_running await pilot.press("ctrl+q") assert not app.is_running @pytest.mark.asyncio async def test_chat_history_can_add_message(temp_db_path: Path): """Test that ChatHistory can display messages.""" from haiku.rag.chat.widgets.chat_history import ChatHistory app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test(): chat_history = app.query_one(ChatHistory) await chat_history.add_message("user", "Hello, how are you?") assert len(chat_history.messages) == 1 assert chat_history.messages[0] == ("user", "Hello, how are you?") await chat_history.add_message("assistant", "I'm doing well, thank you!") assert len(chat_history.messages) == 2 @pytest.mark.asyncio async def test_chat_history_can_add_tool_calls(temp_db_path: Path): """Test that ChatHistory can display inline tool calls.""" from haiku.rag.chat.widgets.chat_history import ChatHistory, ToolCallWidget app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test(): chat_history = app.query_one(ChatHistory) tool_widget = await chat_history.add_tool_call( "tool-1", "search", {"query": "test"} ) assert isinstance(tool_widget, ToolCallWidget) assert tool_widget._completed is False chat_history.mark_tool_complete("tool-1") assert tool_widget._completed is True @pytest.mark.asyncio async def test_chat_history_can_add_citations(temp_db_path: Path): """Test that ChatHistory can display inline citations.""" from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget from haiku.rag.store.models.citation import Citation app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test(): chat_history = app.query_one(ChatHistory) test_citations = [ Citation( index=1, document_id="doc1", chunk_id="chunk1", document_uri="file:///test/doc1.pdf", document_title="Test Document 1", page_numbers=[1, 2], headings=["Section 1"], content="This is some test content from doc 1", ), Citation( index=2, document_id="doc2", chunk_id="chunk2", document_uri="file:///test/doc2.pdf", document_title="Test Document 2", page_numbers=[5], headings=["Section 2", "Subsection"], content="This is test content from doc 2", ), ] await chat_history.add_citations(test_citations) citation_widgets = chat_history.query(CitationWidget) assert len(list(citation_widgets)) == 2 @pytest.mark.asyncio async def test_chat_history_thinking_indicator(temp_db_path: Path): """Test that ChatHistory can show and hide thinking indicator.""" from haiku.rag.chat.widgets.chat_history import ChatHistory, ThinkingWidget app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test() as pilot: chat_history = app.query_one(ChatHistory) await chat_history.show_thinking() thinking = chat_history.query(ThinkingWidget) assert len(list(thinking)) == 1 chat_history.hide_thinking() await pilot.pause() thinking = chat_history.query(ThinkingWidget) assert len(list(thinking)) == 0 @pytest.mark.asyncio async def test_clear_chat_resets_state(temp_db_path: Path): """Test that clearing chat resets state and messages.""" from haiku.rag.chat.widgets.chat_history import ChatHistory app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test() as pilot: chat_history = app.query_one(ChatHistory) await chat_history.add_message("user", "Hello") await chat_history.add_message("assistant", "Hi there") assert len(chat_history.messages) == 2 await app.action_clear_chat() await pilot.pause() assert len(chat_history.messages) == 0 @pytest.mark.asyncio async def test_citation_expand_collapse_with_enter(temp_db_path: Path): """Test that pressing Enter on a focused citation toggles expand/collapse.""" from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget from haiku.rag.store.models.citation import Citation app, mock_client = _make_app(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test() as pilot: chat_history = app.query_one(ChatHistory) test_citation = Citation( index=1, document_id="doc1", chunk_id="chunk1", document_uri="file:///test/doc1.pdf", document_title="Test Document", page_numbers=[1], content="Test content", ) await chat_history.add_citations([test_citation]) citation_widget = chat_history.query_one(CitationWidget) assert citation_widget.collapsed is True citation_widget.focus() await pilot.pause() await pilot.press("enter") await pilot.pause() assert citation_widget.collapsed is False await pilot.press("enter") await pilot.pause() assert citation_widget.collapsed is True @pytest.mark.asyncio async def test_show_citations_renders_from_flat_state(temp_db_path: Path): """Citations in state (flat list[str]) render into the chat history.""" from haiku.rag.chat.app import RAG_STATE_NAMESPACE from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget from haiku.rag.store.models.citation import Citation app, mock_client = _make_app_with_state(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test() as pilot: rag_state = app._toolset.get_namespace(RAG_STATE_NAMESPACE) assert isinstance(rag_state, RAGState) citation = Citation( index=1, document_id="doc1", chunk_id="chunk1", document_uri="file:///test/doc1.pdf", document_title="Test Document", page_numbers=[1], content="Cited content", ) rag_state.citation_index["chunk1"] = citation rag_state.citations.append("chunk1") chat_history = app.query_one(ChatHistory) await app._show_citations_and_programs(chat_history) await pilot.pause() widgets = list(chat_history.query(CitationWidget)) assert len(widgets) == 1 assert widgets[0].citation.chunk_id == "chunk1" @pytest.mark.asyncio async def test_document_filter_updates_rag_state(temp_db_path: Path): """Test that selecting document filters updates RAGState.document_filter.""" from haiku.rag.chat.app import RAG_STATE_NAMESPACE from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal from haiku.rag.tools.filters import build_multi_document_filter app, mock_client = _make_app_with_state(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test(): # Simulate the FilterChanged message selected = ["AI Overview", "ML Basics"] app.on_document_filter_modal_filter_changed( DocumentFilterModal.FilterChanged(selected) ) # RAGState.document_filter should be set rag_state = app._toolset.get_namespace(RAG_STATE_NAMESPACE) assert rag_state is not None expected_filter = build_multi_document_filter(selected) assert rag_state.document_filter == expected_filter # The state snapshot should also reflect the change assert app._state["rag"]["document_filter"] == expected_filter @pytest.mark.asyncio async def test_document_filter_cleared_when_empty(temp_db_path: Path): """Test that clearing all document filters sets document_filter to None.""" from haiku.rag.chat.app import RAG_STATE_NAMESPACE from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal app, mock_client = _make_app_with_state(temp_db_path) with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): async with app.run_test(): # First set a filter app.on_document_filter_modal_filter_changed( DocumentFilterModal.FilterChanged(["AI Overview"]) ) rag_state = app._toolset.get_namespace(RAG_STATE_NAMESPACE) assert rag_state.document_filter is not None # Then clear it app.on_document_filter_modal_filter_changed( DocumentFilterModal.FilterChanged([]) ) assert rag_state.document_filter is None assert app._state["rag"]["document_filter"] is None