""" حفاظ‌های بازبینیِ امنیتیِ پیش از تحویل (دستیار). مدل و 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 = '{"kind":"numberline","min":1' + '0' * 400 + '*2,"max":5}' DEEP_DIAGRAM = '' + '[' * 3000 + ']' * 3000 + '' 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', is_guest=False) 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', is_guest=False) 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'], [])