""" 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)