From 90d179b09a32742ecdc34f82441ebc4cf41e1dbd Mon Sep 17 00:00:00 2001
From: Benjamin Dauvergne <bdauvergne@entrouvert.com>
Date: Wed, 11 May 2022 20:04:44 +0200
Subject: [PATCH 3/9] tests: add fixture decorator for db fixture with global
 scope (#62013)

Creating a django_db fixture which persists between tests is not easy,
this decorator simplify it and completely replace pytest.fixture for
this use case.
---
 tests/conftest.py | 41 ++++++++++++++++++++++++++++++++++++++++-
 tests/utils.py    | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 50ce6341..9a257d7d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -15,6 +15,8 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 
+import contextlib
+import inspect
 import urllib.parse
 from unittest import mock
 
@@ -24,7 +26,7 @@ import pytest
 from django.contrib.auth import get_user_model
 from django.core.cache import cache
 from django.core.management import call_command
-from django.db import connection
+from django.db import connection, transaction
 from django.db.migrations.executor import MigrationExecutor
 
 from authentic2 import hooks as a2_hooks
@@ -528,3 +530,40 @@ def cgu_attribute(db):
         required_on_login=True,
         user_visible=True,
     )
+
+
+@pytest.fixture(scope='session')
+def scoped_db(django_db_setup, django_db_blocker):
+    '''Scoped fixture, use like that to load some models for session/module/class scope:
+
+    @pytest.fixture(scope='module')
+    def myfixture(scoped_db):
+        @scoped_db
+        def f():
+             return Model.objects.create(x=1)
+        yield from f()
+    '''
+
+    @contextlib.contextmanager
+    def scoped_db(func):
+        with django_db_blocker.unblock():
+            with transaction.atomic():
+                try:
+                    if inspect.isgeneratorfunction(func):
+                        generator = func()
+                        value = next(generator)
+                    else:
+                        value = func()
+                    with django_db_blocker.block():
+                        yield value
+                    if inspect.isgeneratorfunction(func):
+                        try:
+                            next(generator)
+                        except StopIteration:
+                            pass
+                        else:
+                            raise RuntimeError(f'{func} yielded more than one time')
+                finally:
+                    transaction.set_rollback(True)
+
+    return scoped_db
diff --git a/tests/utils.py b/tests/utils.py
index 4f21afcb..3034dd29 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -16,11 +16,14 @@
 # authentic2
 
 import base64
+import functools
+import inspect
 import re
 import socket
 import urllib.parse
 from contextlib import closing, contextmanager
 
+import pytest
 from django.contrib.messages.storage.cookie import MessageDecoder, MessageEncoder
 
 try:
@@ -356,3 +359,42 @@ def decode_cookie(data):
     except AttributeError:
         # xxx support legacy decoding?
         return data
+
+
+def scoped_db_fixture(func=None, /, **kwargs):
+    '''Create a db fixture with a scope different than 'function' the default one.'''
+    if func is None:
+        return functools.partial(scoped_db_fixture, **kwargs)
+    assert kwargs.get('scope') in [
+        'session',
+        'module',
+        'class',
+    ], 'scoped_db_fixture is only usable with a non function scope'
+    signature = inspect.signature(func)
+    inner_parameters = []
+    for parameter in signature.parameters:
+        inner_parameters.append(parameter)
+    outer_parameters = ['scoped_db'] if 'scoped_db' not in inner_parameters else []
+    if inner_parameters and inner_parameters[0] == 'self':
+        outer_parameters = ['self'] + outer_parameters + inner_parameters[1:]
+    else:
+        outer_parameters = outer_parameters + inner_parameters
+    # build a new fixture function, inject the scoped_db fixture inside and the old function in a scoped_db context.
+    new_function_def = f'''def f({", ".join(outer_parameters)}):
+    if inspect.isgeneratorfunction(func):
+        def g():
+            yield from func({", ".join(inner_parameters)})
+    else:
+        def g():
+            return func({", ".join(inner_parameters)})
+    with scoped_db(g) as result:
+        yield result'''
+    new_function_bytecode = compile(new_function_def, filename=f'<scoped-db-fixture {func}>', mode='single')
+    global_dict = {'func': func, 'inspect': inspect}
+    eval(new_function_bytecode, global_dict)  # pylint: disable=eval-used
+    new_function = global_dict['f']
+    wrapped_func = functools.wraps(func)(new_function)
+    # prevent original fixture signature to override the new fixture signature
+    # (because of inspect.signature(follow_wrapped=True)) during pytest collection
+    del wrapped_func.__wrapped__
+    return pytest.fixture(**kwargs)(wrapped_func)
-- 
2.35.1