tests/learning_tasks/test_qa.py

from unittest.mock import MagicMock, patch

from langchain import PromptTemplate

from edu_assistant.learning_tasks import QaTask
from edu_assistant.learning_tasks.qa import TEMPLATE_CHAT, TEMPLATE_ONCE


@patch.object(QaTask, "_init_llm")
@patch.object(QaTask, "_build_once_chain")
def test_init_without_knowledge(mocked_build_once_chain, mocked_init_llm):
    task = QaTask(instruction="test")

    assert task._chat_prompt == PromptTemplate.from_template(TEMPLATE_CHAT.format(instruction="test"))
    assert task._once_prompt == PromptTemplate.from_template(TEMPLATE_ONCE.format(instruction="test"))
    assert task._knowledge is None
    mocked_build_once_chain.assert_called_once()


@patch.object(QaTask, "_init_llm")
@patch.object(QaTask, "_build_once_chain")
@patch.object(QaTask, "_create_session_chain")
def test_ask_with_session(mocked_create_session_chain, mocked_build_once_chain, mocked_init_llm):
    mocked_chain = MagicMock(return_value={"response": "ok"})
    mocked_build_once_chain.return_value = mocked_chain
    mocked_create_session_chain.return_value = mocked_chain

    task = QaTask(instruction="test")

    with patch.object(task, "_create_session_id") as mock_create_id:
        mock_create_id.return_value = 123
        result = task.ask("how are you?", session=True)

    mock_create_id.assert_called_once()
    assert "session_id" in result
    assert result["session_id"] == 123
    assert "response" in result
    assert result["response"] == "ok"


@patch.object(QaTask, "_init_llm")
@patch.object(QaTask, "_build_once_chain")
def test_ask_without_session(mocked_build_once_chain, mocked_init_llm):
    mocked_llm = MagicMock()
    mocked_llm.run.return_value = {"result": "ok"}
    mocked_build_once_chain.return_value = mocked_llm
    task = QaTask(instruction="test")

    result = task.ask("how are you?", session=False)

    mocked_build_once_chain.assert_called_once()
    assert "session_id" not in result