tests/unit/helpers/test_message_helper.py

import pytest
from unittest.mock import patch
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from classes.base_models import ChatMessage
from helpers.message_helper import convert_messages, convert_messages_langchain


class TestConvertMessages:
    """Test the convert_messages function"""

    @pytest.fixture
    def sample_messages(self):
        return [
            ChatMessage(role="user", content="Hello"),
            ChatMessage(role="assistant", content="Hi there"),
            ChatMessage(role="user", content="How are you?"),
        ]

    # ==================== Basic Conversion Tests ====================

    def test_convert_messages(self):
        from helpers.message_helper import convert_messages
        from champ.prompts import DEFAULT_SYSTEM_PROMPT_V4

        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_V4.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

    def test_convert_messages_english_no_docs(self, sample_messages):
        """Test basic conversion to English without documents"""
        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System prompt in English"

            result = convert_messages(sample_messages, "en", None)

        assert len(result) == 4  # system + 3 messages
        assert result[0]["role"] == "system"
        assert result[0]["content"] == "System prompt in English"
        assert result[1] == {"role": "user", "content": "Hello"}
        assert result[2] == {"role": "assistant", "content": "Hi there"}
        assert result[3] == {"role": "user", "content": "How are you?"}

    def test_convert_messages_french_no_docs(self, sample_messages):
        """Test basic conversion to French without documents"""
        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "Prompt système en français"

            result = convert_messages(sample_messages, "fr", None)

        assert result[0]["role"] == "system"
        assert result[0]["content"] == "Prompt système en français"

    def test_convert_messages_uses_correct_language_parameter(self, sample_messages):
        """Test that correct language parameter is passed to prompt"""
        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System prompt"

            convert_messages(sample_messages, "en", None)
            mock_prompt.format.assert_called_once_with(language="English")

            mock_prompt.reset_mock()

            convert_messages(sample_messages, "fr", None)
            mock_prompt.format.assert_called_once_with(language="French")

    # ==================== Document Context Tests ====================

    def test_convert_messages_with_documents(self, sample_messages):
        """Test conversion with document context"""
        docs = ["Document 1 content", "Document 2 content"]

        with patch(
            "helpers.message_helper.DEFAULT_SYSTEM_PROMPT_WITH_CONTEXT_V4"
        ) as mock_prompt:
            mock_prompt.format.return_value = "System prompt with context"

            result = convert_messages(sample_messages, "en", docs)

        assert result[0]["role"] == "system"
        assert result[0]["content"] == "System prompt with context"
        mock_prompt.format.assert_called_once_with(context=docs, language="English")

    def test_convert_messages_with_empty_documents(self, sample_messages):
        """Test conversion with empty document list"""
        docs = []

        with patch(
            "helpers.message_helper.DEFAULT_SYSTEM_PROMPT_WITH_CONTEXT_V4"
        ) as mock_prompt:
            mock_prompt.format.return_value = "System prompt with empty context"

            result = convert_messages(sample_messages, "en", docs)

        mock_prompt.format.assert_called_once_with(context=[], language="English")

    def test_convert_messages_none_vs_empty_docs(self, sample_messages):
        """Test that None and empty list are handled differently"""
        with (
            patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_no_context,
            patch(
                "helpers.message_helper.DEFAULT_SYSTEM_PROMPT_WITH_CONTEXT_V4"
            ) as mock_with_context,
        ):
            mock_no_context.format.return_value = "No context"
            mock_with_context.format.return_value = "With context"

            # None should use no-context prompt
            result1 = convert_messages(sample_messages, "en", None)
            assert mock_no_context.format.called
            assert not mock_with_context.format.called

            mock_no_context.reset_mock()
            mock_with_context.reset_mock()

            # Empty list should use with-context prompt
            result2 = convert_messages(sample_messages, "en", [])
            assert not mock_no_context.format.called
            assert mock_with_context.format.called

    # ==================== System Message Filtering Tests ====================

    def test_convert_messages_filters_system_messages(self):
        """Test that system messages in input are filtered out"""
        messages = [
            ChatMessage(role="system", content="Should be filtered"),
            ChatMessage(role="user", content="User message"),
            ChatMessage(role="system", content="Another system message"),
            ChatMessage(role="assistant", content="Assistant message"),
        ]

        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System prompt"

            result = convert_messages(messages, "en", None)

        # Should only have: new system prompt + user + assistant (2 filtered out)
        assert len(result) == 3
        assert result[0]["role"] == "system"
        assert result[1]["role"] == "user"
        assert result[2]["role"] == "assistant"

    def test_convert_messages_only_system_messages(self):
        """Test conversion when only system messages in input"""
        messages = [
            ChatMessage(role="system", content="System 1"),
            ChatMessage(role="system", content="System 2"),
        ]

        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System prompt"

            result = convert_messages(messages, "en", None)

        # Should only have the new system prompt
        assert len(result) == 1
        assert result[0]["role"] == "system"

    # ==================== Message Order Tests ====================

    def test_convert_messages_preserves_order(self):
        """Test that message order is preserved"""
        messages = [
            ChatMessage(role="user", content="1"),
            ChatMessage(role="assistant", content="2"),
            ChatMessage(role="user", content="3"),
            ChatMessage(role="assistant", content="4"),
        ]

        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System"

            result = convert_messages(messages, "en", None)

        assert result[1]["content"] == "1"
        assert result[2]["content"] == "2"
        assert result[3]["content"] == "3"
        assert result[4]["content"] == "4"

    # ==================== Edge Cases ====================

    def test_convert_messages_empty_list(self):
        """Test conversion with empty message list"""
        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System prompt"

            result = convert_messages([], "en", None)

        assert len(result) == 1
        assert result[0]["role"] == "system"

    def test_convert_messages_special_characters(self):
        """Test messages with special characters"""
        messages = [
            ChatMessage(role="user", content="Hello! @#$%^&*() 你好 🎉"),
        ]

        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System"

            result = convert_messages(messages, "en", None)

        assert result[1]["content"] == "Hello! @#$%^&*() 你好 🎉"

    def test_convert_messages_multiline_content(self):
        """Test messages with multiline content"""
        messages = [
            ChatMessage(role="user", content="Line 1\nLine 2\nLine 3"),
        ]

        with patch("helpers.message_helper.DEFAULT_SYSTEM_PROMPT_V4") as mock_prompt:
            mock_prompt.format.return_value = "System"

            result = convert_messages(messages, "en", None)

        assert result[1]["content"] == "Line 1\nLine 2\nLine 3"


class TestConvertMessagesLangchain:
    """Test the convert_messages_langchain function"""

    # ==================== Basic Conversion Tests ====================

    def test_convert_messages_langchain_basic(self):
        """Test basic conversion to LangChain messages"""
        messages = [
            ChatMessage(role="user", content="Hello"),
            ChatMessage(role="assistant", content="Hi"),
            ChatMessage(role="user", content="How are you?"),
        ]

        result = convert_messages_langchain(messages)

        assert len(result) == 3
        assert isinstance(result[0], HumanMessage)
        assert result[0].content == "Hello"
        assert isinstance(result[1], AIMessage)
        assert result[1].content == "Hi"
        assert isinstance(result[2], HumanMessage)
        assert result[2].content == "How are you?"

    def test_convert_messages_langchain_system_message(self):
        """Test conversion of system messages"""
        messages = [
            ChatMessage(role="system", content="You are helpful"),
            ChatMessage(role="user", content="Hello"),
        ]

        result = convert_messages_langchain(messages)

        assert len(result) == 2
        assert isinstance(result[0], SystemMessage)
        assert result[0].content == "You are helpful"
        assert isinstance(result[1], HumanMessage)

    # ==================== History Limiting Tests ====================

    @patch("helpers.message_helper.MAX_HISTORY", 5)
    def test_convert_messages_langchain_respects_max_history(self):
        """Test that only last MAX_HISTORY messages are converted"""
        messages = [ChatMessage(role="user", content=f"Message {i}") for i in range(10)]

        result = convert_messages_langchain(messages)

        assert len(result) == 5
        # Should take the last 5
        assert result[0].content == "Message 5"
        assert result[4].content == "Message 9"

    @patch("helpers.message_helper.MAX_HISTORY", 3)
    def test_convert_messages_langchain_max_history_edge_case(self):
        """Test MAX_HISTORY with exact match"""
        messages = [
            ChatMessage(role="user", content="1"),
            ChatMessage(role="user", content="2"),
            ChatMessage(role="user", content="3"),
        ]

        result = convert_messages_langchain(messages)

        assert len(result) == 3

    @patch("helpers.message_helper.MAX_HISTORY", 10)
    def test_convert_messages_langchain_fewer_than_max(self):
        """Test with fewer messages than MAX_HISTORY"""
        messages = [
            ChatMessage(role="user", content="1"),
            ChatMessage(role="user", content="2"),
        ]

        result = convert_messages_langchain(messages)

        assert len(result) == 2

    # ==================== Message Type Tests ====================

    def test_convert_messages_langchain_all_roles(self):
        """Test conversion of all role types"""
        messages = [
            ChatMessage(role="system", content="System"),
            ChatMessage(role="user", content="User"),
            ChatMessage(role="assistant", content="Assistant"),
        ]

        result = convert_messages_langchain(messages)

        assert isinstance(result[0], SystemMessage)
        assert isinstance(result[1], HumanMessage)
        assert isinstance(result[2], AIMessage)

    def test_convert_messages_langchain_preserves_content(self):
        """Test that content is preserved exactly"""
        content = "Complex content\nWith newlines\n\tAnd tabs"
        messages = [ChatMessage(role="user", content=content)]

        result = convert_messages_langchain(messages)

        assert result[0].content == content

    # ==================== Edge Cases ====================

    def test_convert_messages_langchain_empty_list(self):
        """Test conversion of empty message list"""
        result = convert_messages_langchain([])

        assert result == []

    def test_convert_messages_langchain_special_characters(self):
        """Test messages with special characters"""
        messages = [
            ChatMessage(role="user", content="Hello! 你好 🎉 @#$%"),
        ]

        result = convert_messages_langchain(messages)

        assert result[0].content == "Hello! 你好 🎉 @#$%"

    @patch("helpers.message_helper.MAX_HISTORY", 2)
    def test_convert_messages_langchain_mixed_roles_with_limit(self):
        """Test that role types are preserved when limited by MAX_HISTORY"""
        messages = [
            ChatMessage(role="system", content="System"),
            ChatMessage(role="user", content="User 1"),
            ChatMessage(role="assistant", content="Assistant 1"),
            ChatMessage(role="user", content="User 2"),
            ChatMessage(role="assistant", content="Assistant 2"),
        ]

        result = convert_messages_langchain(messages)

        # Should take last 2
        assert len(result) == 2
        assert isinstance(result[0], HumanMessage)
        assert result[0].content == "User 2"
        assert isinstance(result[1], AIMessage)
        assert result[1].content == "Assistant 2"

    # ==================== Return Type Tests ====================

    def test_convert_messages_langchain_returns_list(self):
        """Test that function returns a list"""
        messages = [ChatMessage(role="user", content="Test")]

        result = convert_messages_langchain(messages)

        assert isinstance(result, list)

    def test_convert_messages_langchain_order_preserved(self):
        """Test that message order is preserved"""
        messages = [
            ChatMessage(role="user", content="1"),
            ChatMessage(role="assistant", content="2"),
            ChatMessage(role="user", content="3"),
        ]

        result = convert_messages_langchain(messages)

        assert result[0].content == "1"
        assert result[1].content == "2"
        assert result[2].content == "3"

    # ==================== Integration Test ====================

    @patch("helpers.message_helper.MAX_HISTORY", 4)
    def test_convert_messages_langchain_realistic_conversation(self):
        """Test with realistic conversation flow"""
        messages = [
            ChatMessage(role="system", content="You are a helpful assistant"),
            ChatMessage(role="user", content="What's the weather?"),
            ChatMessage(role="assistant", content="It's sunny!"),
            ChatMessage(role="user", content="Should I bring an umbrella?"),
            ChatMessage(role="assistant", content="No need today."),
            ChatMessage(role="user", content="Thanks!"),
        ]

        result = convert_messages_langchain(messages)

        # Should take last 4 messages
        assert len(result) == 4
        assert isinstance(result[0], AIMessage)
        assert result[0].content == "It's sunny!"
        assert isinstance(result[-1], HumanMessage)
        assert result[-1].content == "Thanks!"