from langchain_core.messages import AIMessage

from stream_agent import format_message


def test_format_message_with_content():
    message = AIMessage(content="hello")

    assert format_message(message) == "hello"


def test_format_message_with_tool_call():
    message = AIMessage(
        content="",
        tool_calls=[
            {
                "name": "get_demo_price",
                "args": {"product": "молоко", "city": "Казань"},
                "id": "call-1",
            }
        ],
    )

    assert "get_demo_price" in format_message(message)