tests/test_prompt.py

"""Тесты на PromptBuilder.

Покрывают как базовое формирование chat-template, так и опциональную
интеграцию BusinessVocabulary в системное сообщение (раздел 3.6 ВКР).
"""

from src.business.vocabulary import BusinessVocabulary
from src.data.prompt import (
    BASE_SYSTEM_PROMPT,
    SYSTEM_PROMPT,
    build_chat_messages,
    build_system_message,
    build_training_example,
    build_user_message,
)


def test_user_message_contains_parts():
    msg = build_user_message("CREATE TABLE t (id INT);", "Покажи всё")
    assert "Schema:" in msg
    assert "Question:" in msg
    assert "SQL:" in msg
    assert "CREATE TABLE" in msg
    assert "Покажи всё" in msg


def test_chat_messages_have_system_and_user():
    msgs = build_chat_messages("schema", "question")
    assert len(msgs) == 2
    assert msgs[0]["role"] == "system"
    assert msgs[0]["content"] == BASE_SYSTEM_PROMPT
    assert msgs[1]["role"] == "user"


def test_training_example_has_assistant():
    msgs = build_training_example("schema", "question", "SELECT 1")
    assert len(msgs) == 3
    assert msgs[2]["role"] == "assistant"
    assert msgs[2]["content"] == "SELECT 1"


def test_legacy_system_prompt_alias():
    assert SYSTEM_PROMPT == BASE_SYSTEM_PROMPT


def test_system_message_without_vocabulary():
    assert build_system_message(None) == BASE_SYSTEM_PROMPT


def test_system_message_with_empty_vocabulary():
    vocab = BusinessVocabulary.empty()
    assert build_system_message(vocab) == BASE_SYSTEM_PROMPT


def test_system_message_with_terms():
    vocab = BusinessVocabulary(
        company="ООО Ромашка",
        terms={"выручка": "SUM(orders.amount) WHERE orders.status = 'paid'"},
    )
    msg = build_system_message(vocab)
    assert msg.startswith(BASE_SYSTEM_PROMPT)
    assert "ООО Ромашка" in msg
    assert "выручка" in msg
    assert "SUM(orders.amount)" in msg


def test_chat_messages_with_vocabulary_keeps_user_clean():
    vocab = BusinessVocabulary(
        terms={"выручка": "SUM(amount) WHERE status='paid'"},
    )
    msgs = build_chat_messages("schema", "Какая выручка?", vocabulary=vocab)
    assert msgs[0]["role"] == "system"
    assert "SUM(amount)" in msgs[0]["content"]
    assert msgs[1]["role"] == "user"
    assert "SUM(amount)" not in msgs[1]["content"]
    assert "Какая выручка?" in msgs[1]["content"]


def test_training_example_with_vocabulary():
    vocab = BusinessVocabulary(terms={"топ": "ORDER BY x DESC LIMIT 10"})
    msgs = build_training_example(
        "schema", "Топ клиентов", "SELECT 1", vocabulary=vocab
    )
    assert len(msgs) == 3
    assert msgs[0]["role"] == "system"
    assert "ORDER BY x DESC" in msgs[0]["content"]
    assert msgs[2]["role"] == "assistant"