Files
Ai/rag/tests/test_chat_context.py
T
2026-04-24 02:50:27 +03:30

53 lines
1.9 KiB
Python

from unittest.mock import patch
from django.test import SimpleTestCase
from rag.chat import build_chat_context
class ChatContextTests(SimpleTestCase):
@patch("rag.chat.search_with_query")
@patch("rag.chat._rank_text_chunks_by_query")
@patch("rag.chat.chunk_text")
def test_build_chat_context_combines_farm_and_kb_context(
self,
mock_chunk_text,
mock_rank_text_chunks_by_query,
mock_search_with_query,
):
mock_chunk_text.return_value = ["chunk-a", "chunk-b"]
mock_rank_text_chunks_by_query.return_value = ["chunk-b"]
mock_search_with_query.return_value = [
{"text": "kb text 1"},
{"text": "kb text 2"},
]
context = build_chat_context(
query="وضعیت مزرعه چطور است؟",
farm_uuid="farm-123",
farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
)
self.assertIn("[بخش‌های مرتبط بازیابی‌شده از اطلاعات مزرعه]", context)
self.assertIn("chunk-b", context)
self.assertIn("[اطلاعات بازیابی‌شده از پایگاه دانش]", context)
self.assertIn("kb text 1", context)
self.assertIn("kb text 2", context)
@patch("rag.chat.search_with_query", return_value=[])
@patch("rag.chat._rank_text_chunks_by_query", return_value=[])
@patch("rag.chat.chunk_text", return_value=["farm chunk"])
def test_build_chat_context_falls_back_to_full_farm_context(
self,
_mock_chunk_text,
_mock_rank_text_chunks_by_query,
_mock_search_with_query,
):
context = build_chat_context(
query="رطوبت چقدر است؟",
farm_uuid="farm-123",
farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
)
self.assertEqual(context, "")