MathApp/rag/tests.py

"""
تست‌های لایه‌ی RAG. همه‌ی فراخوانی‌های شبکه (OpenAI/GapGPT و Qdrant) mock
می‌شن؛ این تست‌ها بدونِ کلید و بدونِ اینترنت هم سبز می‌مونن.
"""
from types import SimpleNamespace
from unittest.mock import patch

from django.core.cache import cache
from django.test import TestCase, override_settings
from rest_framework.test import APIClient

from accounts.authentication import issue_access_token
from accounts.models import User

from .models import LessonChunk
from .services.clients import reset_clients_cache

RAG_ON = dict(
    OPENAI_API_KEY='test-openai',
    QDRANT_API_KEY='test-qdrant',
    RAG_FORCE_MOCK=False,
    RAG_GRADE_TITLE_OVERRIDES={1: 'ششم دبستان'},
)


def _fake_openai(embedding=None, chat_reply='پاسخِ آزمایشیِ مدل.'):
    dim = 1536
    vec = embedding or [0.01] * dim

    class _Embeddings:
        def create(self, **kwargs):
            return SimpleNamespace(data=[SimpleNamespace(embedding=list(vec))])

    class _Completions:
        def create(self, **kwargs):
            return SimpleNamespace(
                choices=[SimpleNamespace(message=SimpleNamespace(content=chat_reply))]
            )

    return SimpleNamespace(
        embeddings=_Embeddings(),
        chat=SimpleNamespace(completions=_Completions()),
    )


class _FakeQdrant:
    def __init__(self, points=None):
        self._collections = set()
        self._points = points or []
        self.upserted = []

    def get_collections(self):
        return SimpleNamespace(collections=[SimpleNamespace(name=n) for n in self._collections])

    def create_collection(self, collection_name, **kwargs):
        self._collections.add(collection_name)

    def upsert(self, collection_name, points, **kwargs):
        self._collections.add(collection_name)
        self.upserted.extend(points)

    def query_points(self, collection_name, query, limit=4, query_filter=None):
        pts = [SimpleNamespace(score=0.9, payload=p) for p in self._points][:limit]
        return SimpleNamespace(points=pts)


@override_settings(**RAG_ON)
class AskAIRagPathTests(TestCase):
    def setUp(self):
        cache.clear()
        reset_clients_cache()
        self.client = APIClient()
        self.user = User.objects.create(device_id='dev-rag-1', is_guest=True, name='')
        self.token = issue_access_token(self.user.id)

    def _ask(self, body):
        return self.client.post(
            '/api/ai/ask', body, format='json',
            HTTP_AUTHORIZATION=f'Bearer {self.token}',
        )

    def test_ask_with_retrieval_returns_answer_and_used_chunks(self):
        qdrant = _FakeQdrant(points=[
            {'summary': 'چرخه‌ی آب شاملِ تبخیر، تراکم و بارش است.', 'lesson_title': 'آب', 'lesson_number': 3},
            {'summary': 'خورشید آب را گرم می‌کند و بخار می‌شود.', 'lesson_title': 'آب', 'lesson_number': 3},
        ])
        with patch('rag.services.clients.get_openai_client', return_value=_fake_openai(chat_reply='آب تبخیر می‌شود.')), \
             patch('rag.services.clients.get_qdrant_client', return_value=qdrant):
            resp = self._ask({
                'userQuestion': 'چرخه‌ی آب چیست؟',
                'gradeId': 1,
                'subject': 'علوم تجربی',
                'lessonNumber': 3,
            })
        self.assertEqual(resp.status_code, 200, resp.content)
        self.assertEqual(resp.data['answer'], 'آب تبخیر می‌شود.')
        self.assertEqual(len(resp.data['usedChunks']), 2)
        self.assertEqual(resp.data['usedChunks'][0]['title'], 'آب')

    def test_ask_without_subject_skips_retrieval_but_still_answers(self):
        with patch('rag.services.clients.get_openai_client', return_value=_fake_openai(chat_reply='پاسخِ عمومی.')), \
             patch('rag.services.clients.get_qdrant_client', return_value=_FakeQdrant()):
            resp = self._ask({'userQuestion': 'یک سوالِ کلی'})
        self.assertEqual(resp.status_code, 200, resp.content)
        self.assertEqual(resp.data['answer'], 'پاسخِ عمومی.')
        self.assertEqual(resp.data['usedChunks'], [])

    def test_provider_failure_becomes_503_and_releases_slot(self):
        from ai_tutor.models import Conversation, Message

        def boom(*a, **k):
            raise ConnectionError('qdrant down')

        with patch('rag.services.clients.get_openai_client', return_value=_fake_openai()), \
             patch('rag.services.retrieval.search', side_effect=boom):
            resp = self._ask({
                'userQuestion': 'سوال', 'gradeId': 1, 'subject': 'علوم تجربی', 'lessonNumber': 3,
            })
        self.assertEqual(resp.status_code, 503)
        self.assertEqual(Conversation.objects.filter(user_id=self.user.id).count(), 0)
        self.assertEqual(Message.objects.filter(conversation__user_id=self.user.id).count(), 0)


@override_settings(**RAG_ON)
class IngestTextTests(TestCase):
    def setUp(self):
        cache.clear()
        reset_clients_cache()
        self.client = APIClient()

    def test_ingest_text_creates_chunks(self):
        qdrant = _FakeQdrant()
        reply = '{"summary": "خلاصه‌ی این تکه."}'
        with patch('rag.services.clients.get_openai_client', return_value=_fake_openai(chat_reply=reply)), \
             patch('rag.services.clients.get_qdrant_client', return_value=qdrant):
            resp = self.client.post(
                '/api/rag/ingest/text',
                {
                    'lessonId': 12, 'gradeId': 1, 'subject': 'علوم تجربی',
                    'lessonNumber': 3, 'title': 'آب',
                    'text': 'الف' * 900,
                },
                format='json',
                HTTP_X_INTERNAL_KEY='dev-internal-service-key-change-me',
            )
        self.assertEqual(resp.status_code, 200, resp.content)
        self.assertGreater(resp.data['chunksCreated'], 0)
        self.assertEqual(LessonChunk.objects.filter(lesson_id=12).count(), resp.data['chunksCreated'])
        self.assertTrue(qdrant.upserted)
        self.assertEqual(resp.data['collection'], 'علوم تجربی ششم دبستان'.lower())

    def test_ingest_requires_internal_key(self):
        resp = self.client.post(
            '/api/rag/ingest/text',
            {'lessonId': 1, 'subject': 'علوم تجربی', 'lessonNumber': 1, 'text': 'x'},
            format='json',
        )
        # بدونِ X-Internal-Key و بدونِ توکن: DRF بسته به authenticator یا
        # 401 می‌ده یا 403 — هر دو یعنی «رد شد».
        self.assertIn(resp.status_code, (401, 403))


class RagConfigTests(TestCase):
    def test_rag_disabled_without_keys(self):
        from rag.services.config import rag_is_configured

        with override_settings(OPENAI_API_KEY='', QDRANT_API_KEY=''):
            self.assertFalse(rag_is_configured())
        with override_settings(OPENAI_API_KEY='a', QDRANT_API_KEY='b', RAG_FORCE_MOCK=True):
            self.assertFalse(rag_is_configured())
        with override_settings(OPENAI_API_KEY='a', QDRANT_API_KEY='b', RAG_FORCE_MOCK=False):
            self.assertTrue(rag_is_configured())


def _parse_sse(streaming_content):
    """bytes iterator → [(event, data_str), ...]"""
    raw = b''.join(streaming_content).decode('utf-8')
    events = []
    for block in raw.split('\n\n'):
        block = block.strip()
        if not block:
            continue
        ev, data = None, None
        for line in block.splitlines():
            if line.startswith('event: '):
                ev = line[7:]
            elif line.startswith('data: '):
                data = line[6:]
        events.append((ev, data))
    return events


@override_settings(**RAG_ON)
class AskAIStreamTests(TestCase):
    def setUp(self):
        cache.clear()
        reset_clients_cache()
        self.client = APIClient()
        self.user = User.objects.create(device_id='dev-stream-1', is_guest=True, name='')
        self.token = issue_access_token(self.user.id)

    def _ask_stream(self, body):
        return self.client.post(
            '/api/ai/ask', {**body, 'stream': True}, format='json',
            HTTP_AUTHORIZATION=f'Bearer {self.token}',
        )

    def test_stream_emits_chunks_then_done_and_persists(self):
        from ai_tutor.models import Message

        chunks_out, meta = [{'title': 'کسر', 'summary': 'خلاصه‌ی کسر'}], {'collection': 'ریاضی ششم'}
        with patch('rag.services.retrieval.retrieve_context', return_value=(chunks_out, meta)), \
             patch('rag.services.generation.stream_answer', return_value=iter(['برای ', 'جمع ', 'کسرها...'])):
            resp = self._ask_stream({
                'userQuestion': 'جمع کسر؟', 'subject': 'ریاضی', 'grade': 'ششم', 'lessonNumber': 2,
            })
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(resp['content-type'], 'text/event-stream')
        events = _parse_sse(resp.streaming_content)
        kinds = [e for e, _ in events]
        self.assertEqual(kinds.count('chunk'), 3)
        self.assertEqual(kinds[-1], 'done')
        import json as _json
        done = _json.loads(events[-1][1])
        self.assertIn('messageId', done)
        self.assertEqual(done['usedChunks'][0]['title'], 'کسر')
        # پیامِ کاربر و دستیار هر دو ثبت شده و کوتا مصرف شده
        self.assertEqual(Message.objects.filter(conversation__user_id=self.user.id, role='assistant').count(), 1)
        msg = Message.objects.get(conversation__user_id=self.user.id, role='assistant')
        self.assertEqual(msg.text, 'برای جمع کسرها...')

    def test_stream_provider_failure_emits_error_and_releases_slot(self):
        from ai_tutor.models import Conversation, Message

        def boom(*a, **k):
            raise ConnectionError('provider down')

        with patch('rag.services.retrieval.retrieve_context', side_effect=boom):
            resp = self._ask_stream({'userQuestion': 'x', 'subject': 'ریاضی', 'grade': 'ششم', 'lessonNumber': 2})
        events = _parse_sse(resp.streaming_content)
        self.assertEqual(events[-1][0], 'error')
        self.assertEqual(Conversation.objects.filter(user_id=self.user.id).count(), 0)
        self.assertEqual(Message.objects.filter(conversation__user_id=self.user.id).count(), 0)


@override_settings(**RAG_ON, RAG_ANSWER_CACHE_TTL=300)
class AskAIAnswerCacheTests(TestCase):
    def setUp(self):
        cache.clear()
        reset_clients_cache()
        self.client = APIClient()
        self.user = User.objects.create(device_id='dev-cache-1', is_guest=True, name='')
        self.token = issue_access_token(self.user.id)

    def _ask(self, body):
        return self.client.post(
            '/api/ai/ask', body, format='json', HTTP_AUTHORIZATION=f'Bearer {self.token}',
        )

    def test_identical_question_second_time_skips_generation_but_still_persists(self):
        from ai_tutor.models import Message

        qdrant = _FakeQdrant(points=[{'summary': 'خلاصه', 'lesson_title': 'کسر', 'lesson_number': 2}])
        body = {'userQuestion': 'سوال یکسان', 'subject': 'ریاضی', 'grade': 'ششم', 'lessonNumber': 2}

        with patch('rag.services.clients.get_openai_client', return_value=_fake_openai(chat_reply='پاسخ ۱')) as oai, \
             patch('rag.services.clients.get_qdrant_client', return_value=qdrant):
            r1 = self._ask(body)
            calls_after_first = oai.call_count
            r2 = self._ask(body)

        self.assertEqual(r1.status_code, 200)
        self.assertEqual(r2.status_code, 200)
        self.assertEqual(r1.data['answer'], r2.data['answer'])
        # درخواستِ دوم نباید کلاینتِ OpenAI رو دوباره صدا بزنه (کش)
        self.assertEqual(oai.call_count, calls_after_first)
        # ولی پیام و کوتا برای هر دو درخواست ثبت/مصرف شده
        self.assertEqual(Message.objects.filter(conversation__user_id=self.user.id, role='user').count(), 2)
        self.assertGreater(r1.data['remainingQuota'], r2.data['remainingQuota'])