tests/learning_tasks/test_coding_problem.py

from unittest import TestCase
from unittest.mock import MagicMock, patch

from edu_assistant.learning_tasks import CodingProblemAnalysis


class TestCodingProblemAnalysis(TestCase):
    @patch.object(CodingProblemAnalysis, "_init_llm", MagicMock())
    def setUp(self):
        self.analysis = CodingProblemAnalysis()

    @patch.object(
        CodingProblemAnalysis,
        "_build_chain",
        return_value=MagicMock(return_value={"response": "Expected Result"}),
    )
    def test_start_analysis(self, mock_chain):
        problem = MagicMock()
        answer = MagicMock()
        result = self.analysis.start_analysis(problem, answer, "First question")

        self.assertIn("session_id", result)
        self.assertIn("response", result)
        self.assertEqual(result["response"], "Expected Result")
        mock_chain.assert_called_once_with(problem, answer)

    @patch.object(
        CodingProblemAnalysis,
        "_build_chain",
        return_value=MagicMock(return_value={"response": "Expected Result"}),
    )
    def test_ask(self, mock_chain):
        session_id = self.analysis.start_analysis(MagicMock(), MagicMock(), "First question")["session_id"]
        result = self.analysis.ask("New question", session_id)
        self.assertIn("session_id", result)
        self.assertIn("response", result)
        self.assertEqual(result["response"], "Expected Result")
        self.assertEqual(result["session_id"], session_id)
        self.analysis._session_store[session_id].has_calls(2)
        self.analysis._session_store[session_id].assert_called_with({"input": "New question"})