tests/test_main.py

from champ.prompts import CHAMP_SYSTEM_PROMPT, DEFAULT_SYSTEM_PROMPT
from classes.base_models import ChatMessage


class TestMain:
    def test_convert_messages(self):
        from main import convert_and_sanitize_messages

        messages = [
            ChatMessage(role="user", content="Hello"),
            ChatMessage(role="assistant", content="Hi there!"),
        ]
        converted = convert_and_sanitize_messages(messages)
        assert converted == [
            {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there!"},
        ]

    def test_convert_messages_langchain(self):
        from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
        from main import convert_messages_langchain, ChatMessage

        messages = [
            ChatMessage(role="user", content="Hello"),
            ChatMessage(role="assistant", content="Hi there!"),
        ]
        converted = convert_messages_langchain(messages)
        # Check types and content only
        expected_types = [SystemMessage, HumanMessage, AIMessage]
        expected_contents = [CHAMP_SYSTEM_PROMPT, "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