MathApp/config/tests_db_pool.py

"""
poolِ اتصالِ Postgres و رفتارِ زیرِ ازدحام (config/settings.py، config/db.py).

بارِ واقعی با لودتستِ محلی سنجیده شد (Postgresِ ۱۴ با max_connections=60 مثلِ سرور، ۳ worker × ۳۲ نخ،
۱۰۰ کاربرِ هم‌زمان): کدِ قبلی ۲۱٪ِ درخواست‌ها را ۵۰۰ کرد و هر ۵۷ اتصال را گرفت؛ با pool صفر خطا و ۲۴ اتصال.
این تست‌ها تکه‌هایی را نگه می‌دارند که بی‌لودتست هم می‌شود سنجید.
"""
import logging
import os
import subprocess
import sys
from unittest import mock

from django.core.cache import cache
from django.db import OperationalError, connection, transaction
from django.test import RequestFactory, SimpleTestCase, TestCase, TransactionTestCase
from psycopg_pool import PoolTimeout

from accounts.authentication import issue_access_token
from accounts.models import User
from config import health
from config.db import is_pool_timeout, release_db_connection
from config.middleware import DatabaseBusyMiddleware
from content.models import Grade


def _pool_timeout_error():
    """همان شکلی که جنگو PoolTimeout را بالا می‌دهد (wrap_database_errors)."""
    err = OperationalError("couldn't get a connection after 10.00 sec")
    err.__cause__ = PoolTimeout("couldn't get a connection after 10.00 sec")
    return err


class PoolTimeoutDetectionTests(SimpleTestCase):
    def test_only_a_wrapped_pool_timeout_counts(self):
        self.assertTrue(is_pool_timeout(_pool_timeout_error()))
        self.assertFalse(is_pool_timeout(OperationalError('server closed the connection')))
        self.assertFalse(is_pool_timeout(PoolTimeout('raw, not wrapped by django')))
        self.assertFalse(is_pool_timeout(ValueError()))


class ServerBusyResponseTests(TestCase):
    def setUp(self):
        cache.clear()
        self.user = User.objects.create(phone='09120007711')
        self.auth = {'HTTP_AUTHORIZATION': f'Bearer {issue_access_token(self.user.id)}'}

    def test_drf_view_answers_503_with_retry_after_not_500(self):
        # هندلرِ خطای DRF (config/exceptions.py) هر خطای ناشناخته را خودش ۵۰۰ می‌کند و به میان‌افزار
        # نمی‌رسد؛ اولین نسخه‌ی این تغییر فقط میان‌افزار را داشت و در لودتست همه ۵۰۰ ماند.
        with mock.patch('content.views.HomeView.get', side_effect=_pool_timeout_error()), \
                self.assertLogs('config.db', level='ERROR') as logs:
            response = self.client.get('/api/home', **self.auth)
        self.assertEqual(response.status_code, 503)
        self.assertEqual(response['Retry-After'], '5')
        self.assertEqual(response.json()['code'], 'server_busy')
        self.assertIn('db_pool_exhausted', logs.output[0])

    def test_other_database_errors_stay_500(self):
        with mock.patch('content.views.HomeView.get', side_effect=OperationalError('db down')), \
                self.assertLogs('config.exceptions', level='ERROR'):
            response = self.client.get('/api/home', **self.auth)
        self.assertEqual(response.status_code, 500)

    def test_middleware_covers_non_drf_views(self):
        request = RequestFactory().get('/admin/')
        middleware = DatabaseBusyMiddleware(lambda r: None)
        with self.assertLogs('config.db', level='ERROR'):
            response = middleware.process_exception(request, _pool_timeout_error())
        self.assertEqual(response.status_code, 503)
        self.assertEqual(response['Retry-After'], '5')
        self.assertIsNone(middleware.process_exception(request, OperationalError('db down')))


class ReleaseConnectionTests(TransactionTestCase):
    def test_outside_a_transaction_the_connection_goes_back(self):
        User.objects.count()
        self.assertIsNotNone(connection.connection)
        release_db_connection()
        self.assertIsNone(connection.connection)
        self.assertEqual(User.objects.count(), 0)  # کوئریِ بعدی خودش یکی تازه می‌گیرد

    def test_inside_atomic_it_keeps_the_connection_and_warns(self):
        with transaction.atomic():
            User.objects.count()
            with self.assertLogs('config.db', level='WARNING'):
                release_db_connection()
            self.assertIsNotNone(connection.connection)
            User.objects.create(phone='09120007712')  # تراکنش سالم ماند
        self.assertEqual(User.objects.count(), 1)


class ReleaseUnderTestCaseIsSilentTests(TestCase):
    def test_testcase_transaction_is_not_reported_as_a_bug(self):
        logger = logging.getLogger('config.db')
        with mock.patch.object(logger, 'warning') as warning:
            release_db_connection()
        warning.assert_not_called()


class AskReleasesTheConnectionBeforeTheModelTests(TransactionTestCase):
    """دستیار تا ۴۰ ثانیه منتظرِ Qdrant و مدل است؛ در این مدت نباید اتصالی از pool دستش باشد."""

    def setUp(self):
        cache.clear()
        grade = Grade.objects.create(title='پایه نهم', icon='', color='')
        self.user = User.objects.create(phone='09120007713', grade_id=grade.id)
        self.seen = {}

    def _ask(self, **extra):
        return self.client.post(
            '/api/ai/ask', {'userQuestion': 'توان چیست؟', **extra}, content_type='application/json',
            HTTP_AUTHORIZATION=f'Bearer {issue_access_token(self.user.id)}',
        )

    def _retrieve(self, *args, **kwargs):
        self.seen['retrieve'] = connection.connection
        return [{'title': 'فصل', 'summary': 'متن'}], {'gradeTitle': 'نهم'}

    def test_json_path(self):
        def generate(*args, **kwargs):
            self.seen['generate'] = connection.connection
            return 'جواب', [], {}

        with mock.patch('rag.services.config.rag_is_configured', return_value=True), \
                mock.patch('rag.services.retrieval.retrieve_context', side_effect=self._retrieve), \
                mock.patch('rag.services.generation.generate_answer', side_effect=generate):
            response = self._ask()
        self.assertEqual(response.status_code, 200)
        self.assertIsNone(self.seen['retrieve'])
        self.assertIsNone(self.seen['generate'])

    def test_stream_path(self):
        def stream(*args, **kwargs):
            self.seen['stream'] = connection.connection
            yield 'جواب'

        with mock.patch('rag.services.config.rag_is_configured', return_value=True), \
                mock.patch('rag.services.retrieval.retrieve_context', side_effect=self._retrieve), \
                mock.patch('rag.services.generation.stream_answer', side_effect=stream):
            response = self._ask(stream=True)
            body = b''.join(response.streaming_content).decode()
        self.assertIn('event: done', body)
        self.assertIsNone(self.seen['retrieve'])
        self.assertIsNone(self.seen['stream'])


class PoolBudgetHealthTests(SimpleTestCase):
    def test_over_budget_is_a_warning(self):
        warnings = health._db_pool_warnings(
            {'state': 'over_budget', 'workers': 6, 'poolMaxSize': 10, 'budget': 66, 'maxConnections': 57})
        self.assertEqual(len(warnings), 1)
        self.assertIn('66', warnings[0])
        self.assertEqual(health._db_pool_warnings({'state': 'ok'}), [])
        self.assertEqual(health._db_pool_warnings({'state': 'not_postgres'}), [])


class PoolBudgetOnPostgresTests(TestCase):
    def test_reports_the_real_limits(self):
        if connection.vendor != 'postgresql':
            self.skipTest('pool فقط روی Postgres است')
        with mock.patch.dict(os.environ, {'GUNICORN_WORKERS': '4'}):
            state = health._db_pool()
        self.assertEqual(state['workers'], 4)
        self.assertEqual(state['budget'], 4 * state['poolMaxSize'] + health.DB_CONNECTION_RESERVE)
        self.assertGreater(state['maxConnections'], 0)
        self.assertIn(state['state'], ('ok', 'over_budget'))


class PoolSettingsFailLoudlyTests(SimpleTestCase):
    """تنظیمِ غلط باید سرویس را بالا نیاورد، نه این‌که سرِ اولین کوئری همه‌ی درخواست‌ها ۵۰۰ شوند."""

    def _load_settings(self, **env):
        base = {k: v for k, v in os.environ.items() if not k.startswith('DB_')}
        base.update({'DB_NAME': 'x', 'DB_USER': 'x', 'DB_PASSWORD': 'x', **env})
        return subprocess.run(
            [sys.executable, '-c', 'import config.settings'],
            env=base, capture_output=True, text=True, cwd=os.path.dirname(os.path.dirname(__file__)),
        )

    def test_good_settings_load(self):
        self.assertEqual(self._load_settings().returncode, 0)

    def test_old_conn_max_age_is_refused(self):
        result = self._load_settings(DB_CONN_MAX_AGE='60')
        self.assertNotEqual(result.returncode, 0)
        self.assertIn('DB_CONN_MAX_AGE', result.stderr)

    def test_min_above_max_is_refused(self):
        result = self._load_settings(DB_POOL_MIN_SIZE='3', DB_POOL_MAX_SIZE='2')
        self.assertNotEqual(result.returncode, 0)
        self.assertIn('DB_POOL_MIN_SIZE', result.stderr)