MathApp/ai_tutor/tests_security_review.py

"""
حفاظ‌های بازبینیِ امنیتیِ پیش از تحویل (دستیار). مدل و Qdrant همه ساختگی‌اند.
"""
import json
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 ai_tutor.models import Message
from rag.services.clients import reset_clients_cache

RAG_ON = dict(OPENAI_API_KEY='t', QDRANT_API_KEY='t', RAG_FORCE_MOCK=False, AI_DAILY_QUOTA_FREE=3)

ANSWER = '۱) دو به‌علاوه‌ی دو می‌شود چهار.\n✅ جواب: ۴\n'
OVERFLOW_DIAGRAM = '<diagram>{"kind":"numberline","min":1' + '0' * 400 + '*2,"max":5}</diagram>'
DEEP_DIAGRAM = '<diagram>' + '[' * 3000 + ']' * 3000 + '</diagram>'


class FakeOpenAI:
    def __init__(self, reply):
        self.reply = reply
        self.calls = 0
        self.chat = SimpleNamespace(completions=SimpleNamespace(create=self._create))
        self.embeddings = SimpleNamespace(
            create=lambda **k: SimpleNamespace(data=[SimpleNamespace(embedding=[0.01] * 1536)]))

    def _create(self, **kw):
        self.calls += 1
        if kw.get('stream'):
            parts = [self.reply[i:i + 25] for i in range(0, len(self.reply), 25)]
            events = [SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=p))], usage=None)
                      for p in parts]
            events.append(SimpleNamespace(choices=[], usage=SimpleNamespace(prompt_tokens=900, completion_tokens=300)))
            return iter(events)
        return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content=self.reply))],
                               usage=SimpleNamespace(prompt_tokens=900, completion_tokens=300))


def sse(resp):
    out = []
    for block in b''.join(resp.streaming_content).decode().split('\n\n'):
        ev = data = None
        for line in block.strip().splitlines():
            if line.startswith('event: '):
                ev = line[7:]
            elif line.startswith('data: '):
                data = json.loads(line[6:])
        if ev:
            out.append((ev, data))
    return out


@override_settings(**RAG_ON)
class _Base(TestCase):
    def setUp(self):
        cache.clear()
        reset_clients_cache()
        self.user = User.objects.create(phone='09120009901')
        self.c = APIClient()
        self.c.credentials(HTTP_AUTHORIZATION=f'Bearer {issue_access_token(self.user.id)}')

    def ask(self, q, **extra):
        return self.c.post('/api/ai/ask', {'userQuestion': q, **extra}, format='json')

    def remaining(self):
        return self.c.get('/api/ai/quota').data['remainingQuota']


class BrokenDiagramDoesNotMakeAnswersFreeTests(_Base):
    """پارسرِ شکل OverflowError/RecursionError را بالا می‌داد و جوابِ استریم‌شده رایگان و بی‌ثبت می‌ماند."""

    def test_stream_and_json_answers_are_charged(self):
        for diagram in (OVERFLOW_DIAGRAM, DEEP_DIAGRAM):
            with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI(ANSWER + diagram)):
                events = sse(self.ask(f'دو به‌علاوه‌ی دو؟ {len(diagram)}', stream=True))
            self.assertEqual(events[-1][0], 'done')
            self.assertIsNone(events[-1][1]['diagram'])
        with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI(ANSWER + OVERFLOW_DIAGRAM)):
            self.assertEqual(self.ask('دو به‌علاوه‌ی دو؟').status_code, 200)
        self.assertEqual(self.remaining(), 0)
        self.assertEqual(Message.objects.filter(role='assistant').count(), 3)


@override_settings(RAG_ANSWER_CACHE_TTL=300)
class LateOffTopicMarkerTests(_Base):
    def test_rest_is_not_streamed_stored_or_cached(self):
        from rag.services.guardrails import OFF_TOPIC_MARKER, REFUSAL

        reply = ('۱) این یک جمله‌ی ریاضیِ بلند برای عبور از چهل نویسه است.\n'
                 f'{OFF_TOPIC_MARKER} متنِ غیرِ ریاضی: یک انشای کامل درباره‌ی فوتبال.\n')
        fake = FakeOpenAI(reply)
        q = 'سوال درباره‌ی چیزی'
        with patch('rag.services.clients.get_openai_client', return_value=fake):
            streamed = ''.join(d['delta'] for e, d in sse(self.ask(q, stream=True)) if e == 'chunk')
            other = User.objects.create(phone='09120009902')
            c2 = APIClient()
            c2.credentials(HTTP_AUTHORIZATION=f'Bearer {issue_access_token(other.id)}')
            second = c2.post('/api/ai/ask', {'userQuestion': q}, format='json')
        self.assertNotIn('فوتبال', streamed)
        self.assertIn(REFUSAL, streamed)
        self.assertEqual(Message.objects.filter(role='assistant').first().text, REFUSAL)
        self.assertNotIn('فوتبال', second.data['answer'])
        self.assertEqual(fake.calls, 2)  # جوابِ ردشده در کش نرفت


class DisconnectMidStreamTests(_Base):
    def test_partial_answer_is_saved(self):
        with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI(ANSWER * 10)):
            resp = self.ask('سوال', stream=True)
            it = iter(resp.streaming_content)
            next(it)
            next(it)
            resp._iterator.close()  # همان GeneratorExit که بستنِ صفحه می‌سازد
        saved = Message.objects.filter(role='assistant')
        self.assertEqual(saved.count(), 1)
        self.assertTrue(saved.first().text.endswith('…'))
        self.assertEqual(self.remaining(), 2)


class UnknownCollectionIsNotSearchedTests(_Base):
    def test_client_chosen_collection_is_ignored(self):
        seen = []

        def fake_search(collection, vec, **kw):
            seen.append(collection)
            return [{'score': 0.9, 'payload': {'title': 'doc', 'text': 'PAYLOAD-OF-UNRELATED-COLLECTION'}}]

        with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI('پاسخ')), \
                patch('rag.services.retrieval.search', side_effect=fake_search):
            r = self.ask('هر چیزی', subject='any', grade='other collection')
        self.assertEqual(r.status_code, 200)
        self.assertEqual(seen, [])
        self.assertEqual(r.data['usedChunks'], [])