tests/test_main.py

from classes.base_models import ChatMessage


class TestMain:
    def test_convert_messages(self):
        from helpers.message_helper import convert_messages
        from champ.prompts import DEFAULT_SYSTEM_PROMPT_V3

        messages = [
            ChatMessage(role="user", content="Hello"),
            ChatMessage(role="assistant", content="Hi there!"),
        ]
        lang = "en"
        converted = convert_messages(messages, lang)
        assert converted == [
            {
                "role": "system",
                "content": DEFAULT_SYSTEM_PROMPT_V3.format(language="English"),
            },
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there!"},
        ]

    def test_convert_messages_langchain(self):
        from helpers.message_helper import convert_messages_langchain
        from langchain_core.messages import HumanMessage, AIMessage

        messages = [
            ChatMessage(role="user", content="Hello"),
            ChatMessage(role="assistant", content="Hi there!"),
        ]
        converted = convert_messages_langchain(messages)
        # Check types and content only
        expected_types = [HumanMessage, AIMessage]
        expected_contents = ["Hello", "Hi there!"]

        for msg, expected_type, expected_content in zip(
            converted, expected_types, expected_contents
        ):
            assert isinstance(msg, expected_type)
            assert msg.content == expected_content