from io import StringIO from pathlib import Path from typing import Any from unittest.mock import MagicMock import pytest from pytest import MonkeyPatch from rich.console import Console from typer.testing import CliRunner import llm_chat from llm_chat.chat import ChatProtocol from llm_chat.cli import app from llm_chat.models import Message, Role runner = CliRunner() class ChatFake: """Fake chat class for testing.""" args: dict[str, Any] received_messages: list[str] def __init__(self) -> None: self.args = {} self.received_messages = [] def _set_args(self, **kwargs: Any) -> None: self.args = kwargs @property def cost(self) -> float: """Get the cost of the conversation.""" return 0.0 def save(self) -> None: """Dummy save method.""" pass def send_message(self, message: str) -> str: """Echo the received message.""" self.received_messages.append(message) return message def test_chat(monkeypatch: MonkeyPatch) -> None: chat_fake = ChatFake() output = StringIO() console = Console(file=output) def mock_get_chat(**_: Any) -> ChatProtocol: return chat_fake def mock_get_console() -> Console: return console mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat) monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console) monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input) result = runner.invoke(app) assert result.exit_code == 0 assert chat_fake.received_messages == ["Hello"] @pytest.mark.parametrize("argument", ["--context", "-c"], ids=["--context", "-c"]) def test_chat_with_context( argument: str, monkeypatch: MonkeyPatch, tmp_path: Path ) -> None: context = "Hello, world!" tmp_file = tmp_path / "context.txt" tmp_file.write_text(context) chat_fake = ChatFake() output = StringIO() console = Console(file=output) def mock_get_chat(**kwargs: Any) -> ChatProtocol: chat_fake._set_args(**kwargs) return chat_fake def mock_get_console() -> Console: return console mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat) monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console) monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input) result = runner.invoke(app, [argument, str(tmp_file)]) assert result.exit_code == 0 assert chat_fake.received_messages == ["Hello"] assert "context" in chat_fake.args assert chat_fake.args["context"] == [Message(role=Role.SYSTEM, content=context)]