Ai/rag/tests/test_chat_context.py

from unittest.mock import patch

from django.test import SimpleTestCase

from rag.chat import build_rag_context


class ChatContextTests(SimpleTestCase):
    @patch("rag.chat.search_with_texts")
    @patch("rag.chat.chunk_text")
    def test_build_rag_context_includes_full_farm_and_kb_results(
        self,
        mock_chunk_text,
        mock_search_with_texts,
    ):
        mock_chunk_text.return_value = ["farm chunk 1", "farm chunk 2"]
        mock_search_with_texts.return_value = [
            {"id": "kb-1", "score": 0.8, "text": "kb text 1", "metadata": {}},
            {"id": "kb-2", "score": 0.7, "text": "kb text 2", "metadata": {}},
        ]

        context = build_rag_context(
            query="وضعیت مزرعه چطور است؟",
            sensor_uuid="farm-123",
            service_id="chat",
            farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
        )

        self.assertIn("[اطلاعات کامل مزرعه]", context)
        self.assertIn("soil_moisture", context)
        self.assertIn("[متن‌های مرجع]", context)
        self.assertIn("kb text 1", context)
        self.assertIn("kb text 2", context)
        mock_search_with_texts.assert_called_once()
        sent_texts = mock_search_with_texts.call_args.kwargs["texts"]
        self.assertEqual(sent_texts[0], "وضعیت مزرعه چطور است؟")
        self.assertIn("farm chunk 1", sent_texts)
        self.assertIn("farm chunk 2", sent_texts)

    @patch("rag.chat.search_with_texts", return_value=[])
    @patch("rag.chat.chunk_text", return_value=["farm chunk"])
    def test_build_rag_context_returns_full_farm_when_kb_empty(
        self,
        _mock_chunk_text,
        _mock_search_with_texts,
    ):
        context = build_rag_context(
            query="رطوبت چقدر است؟",
            sensor_uuid="farm-123",
            service_id="chat",
            farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
        )

        self.assertIn("[اطلاعات کامل مزرعه]", context)
        self.assertIn("soil_moisture", context)