Add support for rollback emulation/serialized rollback

bluetech · bluetech · commit 66a21e43f153 · 2021-12-01T15:26:05.000+02:00
Thanks to Aymeric Augustin, Daniel Hahler and Fábio C. Barrionuevo da Luz for previous work on this feature. Fix pytest-dev#329. Closes pytest-dev#353. Closes pytest-dev#721. Closes pytest-dev#919. Closes pytest-dev#956.
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -4,11 +4,18 @@ Changelog
 unreleased
 ----------
 
+Improvements
+^^^^^^^^^^^^
+
+* Add support for :ref:`rollback emulation/serialized rollback
+  <test-case-serialized-rollback>`. The :func:`pytest.mark.django_db` marker
+  has a new ``serialized_rollback`` option, and a
+  :fixture:`django_db_serialized_rollback` fixture is added.
+
 Bugfixes
 ^^^^^^^^
 
-* Fix :fixture:`live_server` when using an in-memory SQLite database on
-  Django >= 3.0.
+* Fix :fixture:`live_server` when using an in-memory SQLite database.
 
 
 v4.4.0 (2021-06-06)
diff --git a/docs/helpers.rst b/docs/helpers.rst
@@ -73,15 +73,27 @@ dynamically in a hook or fixture.
     For details see :py:attr:`django.test.TransactionTestCase.databases` and
     :py:attr:`django.test.TestCase.databases`.
 
+  :type serialized_rollback: bool
+  :param serialized_rollback:
+    The ``serialized_rollback`` argument enables :ref:`rollback emulation
+    <test-case-serialized-rollback>`.  After a transactional test (or any test
+    using a database backend which doesn't support transactions) runs, the
+    database is flushed, destroying data created in data migrations.  Setting
+    ``serialized_rollback=True`` tells Django to serialize the database content
+    during setup, and restore it during teardown.
+
+    Note that this will slow down that test suite by approximately 3x.
+
 .. note::
 
   If you want access to the Django database inside a *fixture*, this marker may
   or may not help even if the function requesting your fixture has this marker
-  applied, depending on pytest's fixture execution order. To access the
-  database in a fixture, it is recommended that the fixture explicitly request
-  one of the :fixture:`db`, :fixture:`transactional_db` or
-  :fixture:`django_db_reset_sequences` fixtures. See below for a description of
-  them.
+  applied, depending on pytest's fixture execution order. To access the database
+  in a fixture, it is recommended that the fixture explicitly request one of the
+  :fixture:`db`, :fixture:`transactional_db`,
+  :fixture:`django_db_reset_sequences` or
+  :fixture:`django_db_serialized_rollback` fixtures. See below for a description
+  of them.
 
 .. note:: Automatic usage with ``django.test.TestCase``.
 
@@ -331,6 +343,17 @@ fixtures which need database access themselves. A test function should normally
 use the :func:`pytest.mark.django_db` mark with ``transaction=True`` and
 ``reset_sequences=True``.
 
+.. fixture:: django_db_serialized_rollback
+
+``django_db_serialized_rollback``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This fixture triggers :ref:`rollback emulation <test-case-serialized-rollback>`.
+This is only required for fixtures which need to enforce this behavior.  A test
+function should normally use :func:`pytest.mark.django_db` with
+``serialized_rollback=True`` (and most likely also ``transaction=True``) to
+request this behavior.
+
 .. fixture:: live_server
 
 ``live_server``
@@ -342,17 +365,23 @@ or by requesting it's string value: ``str(live_server)``.  You can
 also directly concatenate a string to form a URL: ``live_server +
 '/foo'``.
 
+Since the live server and the tests run in different threads, they
+cannot share a database transaction. For this reason, ``live_server``
+depends on the ``transactional_db`` fixture. If tests depend on data
+created in data migrations, you should add the
+``django_db_serialized_rollback`` fixture.
+
 .. note:: Combining database access fixtures.
 
   When using multiple database fixtures together, only one of them is
   used.  Their order of precedence is as follows (the last one wins):
 
   * ``db``
   * ``transactional_db``
-  * ``django_db_reset_sequences``
 
-  In addition, using ``live_server`` will also trigger transactional
-  database access, if not specified.
+  In addition, using ``live_server`` or ``django_db_reset_sequences`` will also
+  trigger transactional database access, and ``django_db_serialized_rollback``
+  regular database access, if not specified.
 
 .. fixture:: settings
 
diff --git a/pytest_django/fixtures.py b/pytest_django/fixtures.py
@@ -20,14 +20,16 @@
     import django
 
     _DjangoDbDatabases = Optional[Union["Literal['__all__']", Iterable[str]]]
-    _DjangoDb = Tuple[bool, bool, _DjangoDbDatabases]
+    # transaction, reset_sequences, databases, serialized_rollback
+    _DjangoDb = Tuple[bool, bool, _DjangoDbDatabases, bool]
 
 
 __all__ = [
     "django_db_setup",
     "db",
     "transactional_db",
     "django_db_reset_sequences",
+    "django_db_serialized_rollback",
     "admin_user",
     "django_user_model",
     "django_username_field",
@@ -151,9 +153,19 @@ def _django_db_helper(
 
     marker = request.node.get_closest_marker("django_db")
     if marker:
-        transactional, reset_sequences, databases = validate_django_db(marker)
+        (
+            transactional,
+            reset_sequences,
+            databases,
+            serialized_rollback,
+        ) = validate_django_db(marker)
     else:
-        transactional, reset_sequences, databases = False, False, None
+        (
+            transactional,
+            reset_sequences,
+            databases,
+            serialized_rollback,
+        ) = False, False, None, False
 
     transactional = transactional or (
         "transactional_db" in request.fixturenames
@@ -162,6 +174,9 @@ def _django_db_helper(
     reset_sequences = reset_sequences or (
         "django_db_reset_sequences" in request.fixturenames
     )
+    serialized_rollback = serialized_rollback or (
+        "django_db_serialized_rollback" in request.fixturenames
+    )
 
     django_db_blocker.unblock()
     request.addfinalizer(django_db_blocker.restore)
@@ -175,10 +190,12 @@ def _django_db_helper(
         test_case_class = django.test.TestCase
 
     _reset_sequences = reset_sequences
+    _serialized_rollback = serialized_rollback
     _databases = databases
 
     class PytestDjangoTestCase(test_case_class):  # type: ignore[misc,valid-type]
         reset_sequences = _reset_sequences
+        serialized_rollback = _serialized_rollback
         if _databases is not None:
             databases = _databases
 
@@ -196,18 +213,20 @@ def validate_django_db(marker) -> "_DjangoDb":
     """Validate the django_db marker.
 
     It checks the signature and creates the ``transaction``,
-    ``reset_sequences`` and ``databases`` attributes on the marker
-    which will have the correct values.
+    ``reset_sequences``, ``databases`` and ``serialized_rollback`` attributes on
+    the marker which will have the correct values.
 
-    A sequence reset is only allowed when combined with a transaction.
+    Sequence reset and serialized_rollback are only allowed when combined with
+    transaction.
     """
 
     def apifun(
         transaction: bool = False,
         reset_sequences: bool = False,
         databases: "_DjangoDbDatabases" = None,
+        serialized_rollback: bool = False,
     ) -> "_DjangoDb":
-        return transaction, reset_sequences, databases
+        return transaction, reset_sequences, databases, serialized_rollback
 
     return apifun(*marker.args, **marker.kwargs)
 
@@ -303,6 +322,27 @@ def django_db_reset_sequences(
     # is requested.
 
 
+@pytest.fixture(scope="function")
+def django_db_serialized_rollback(
+    _django_db_helper: None,
+    db: None,
+) -> None:
+    """Require a test database with serialized rollbacks.
+
+    This requests the ``db`` fixture, and additionally performs rollback
+    emulation - serializes the database contents during setup and restores
+    it during teardown.
+
+    This fixture may be useful for transactional tests, so is usually combined
+    with ``transactional_db``, but can also be useful on databases which do not
+    support transactions.
+
+    Note that this will slow down that test suite by approximately 3x.
+    """
+    # The `_django_db_helper` fixture checks if `django_db_serialized_rollback`
+    # is requested.
+
+
 @pytest.fixture()
 def client() -> "django.test.client.Client":
     """A Django test client instance."""
diff --git a/pytest_django/plugin.py b/pytest_django/plugin.py
@@ -33,6 +33,7 @@
 from .fixtures import django_db_modify_db_settings_tox_suffix  # noqa
 from .fixtures import django_db_modify_db_settings_xdist_suffix  # noqa
 from .fixtures import django_db_reset_sequences  # noqa
+from .fixtures import django_db_serialized_rollback  # noqa
 from .fixtures import django_db_setup  # noqa
 from .fixtures import django_db_use_migrations  # noqa
 from .fixtures import django_user_model  # noqa
@@ -265,14 +266,17 @@ def pytest_load_initial_conftests(
     # Register the marks
     early_config.addinivalue_line(
         "markers",
-        "django_db(transaction=False, reset_sequences=False, databases=None): "
+        "django_db(transaction=False, reset_sequences=False, databases=None, "
+        "serialized_rollback=False): "
         "Mark the test as using the Django test database.  "
         "The *transaction* argument allows you to use real transactions "
         "in the test like Django's TransactionTestCase.  "
         "The *reset_sequences* argument resets database sequences before "
         "the test.  "
         "The *databases* argument sets which database aliases the test "
-        "uses (by default, only 'default'). Use '__all__' for all databases.",
+        "uses (by default, only 'default'). Use '__all__' for all databases.  "
+        "The *serialized_rollback* argument enables rollback emulation for "
+        "the test.",
     )
     early_config.addinivalue_line(
         "markers",
@@ -387,7 +391,12 @@ def get_order_number(test: pytest.Item) -> int:
         else:
             marker_db = test.get_closest_marker('django_db')
             if marker_db:
-                transaction, reset_sequences, databases = validate_django_db(marker_db)
+                (
+                    transaction,
+                    reset_sequences,
+                    databases,
+                    serialized_rollback,
+                ) = validate_django_db(marker_db)
                 uses_db = True
                 transactional = transaction or reset_sequences
             else:
diff --git a/tests/test_database.py b/tests/test_database.py
@@ -48,14 +48,23 @@ def non_zero_sequences_counter(db: None) -> None:
 class TestDatabaseFixtures:
     """Tests for the different database fixtures."""
 
-    @pytest.fixture(params=["db", "transactional_db", "django_db_reset_sequences"])
+    @pytest.fixture(params=[
+        "db",
+        "transactional_db",
+        "django_db_reset_sequences",
+        "django_db_serialized_rollback",
+    ])
     def all_dbs(self, request) -> None:
         if request.param == "django_db_reset_sequences":
             return request.getfixturevalue("django_db_reset_sequences")
         elif request.param == "transactional_db":
             return request.getfixturevalue("transactional_db")
         elif request.param == "db":
             return request.getfixturevalue("db")
+        elif request.param == "django_db_serialized_rollback":
+            return request.getfixturevalue("django_db_serialized_rollback")
+        else:
+            assert False  # pragma: no cover
 
     def test_access(self, all_dbs: None) -> None:
         Item.objects.create(name="spam")
@@ -113,6 +122,51 @@ def test_django_db_reset_sequences_requested(
             ["*test_django_db_reset_sequences_requested PASSED*"]
         )
 
+    def test_serialized_rollback(self, db: None, django_testdir) -> None:
+        django_testdir.create_app_file(
+            """
+            from django.db import migrations
+
+            def load_data(apps, schema_editor):
+                Item = apps.get_model("app", "Item")
+                Item.objects.create(name="loaded-in-migration")
+
+            class Migration(migrations.Migration):
+                dependencies = [
+                    ("app", "0001_initial"),
+                ]
+
+                operations = [
+                    migrations.RunPython(load_data),
+                ]
+            """,
+            "migrations/0002_data_migration.py",
+        )
+
+        django_testdir.create_test_module(
+            """
+            import pytest
+            from .app.models import Item
+
+            @pytest.mark.django_db(transaction=True, serialized_rollback=True)
+            def test_serialized_rollback_1():
+                assert Item.objects.filter(name="loaded-in-migration").exists()
+
+            @pytest.mark.django_db(transaction=True)
+            def test_serialized_rollback_2(django_db_serialized_rollback):
+                assert Item.objects.filter(name="loaded-in-migration").exists()
+                Item.objects.create(name="test2")
+
+            @pytest.mark.django_db(transaction=True, serialized_rollback=True)
+            def test_serialized_rollback_3():
+                assert Item.objects.filter(name="loaded-in-migration").exists()
+                assert not Item.objects.filter(name="test2").exists()
+            """
+        )
+
+        result = django_testdir.runpytest_subprocess("-v")
+        assert result.ret == 0
+
     @pytest.fixture
     def mydb(self, all_dbs: None) -> None:
         # This fixture must be able to access the database
@@ -160,6 +214,10 @@ def fixture_with_transdb(self, transactional_db: None) -> None:
     def fixture_with_reset_sequences(self, django_db_reset_sequences: None) -> None:
         Item.objects.create(name="spam")
 
+    @pytest.fixture
+    def fixture_with_serialized_rollback(self, django_db_serialized_rollback: None) -> None:
+        Item.objects.create(name="ham")
+
     def test_trans(self, fixture_with_transdb: None) -> None:
         pass
 
@@ -180,6 +238,16 @@ def test_reset_sequences(
     ) -> None:
         pass
 
+    # The test works when transactions are not supported, but it interacts
+    # badly with other tests.
+    @pytest.mark.skipif('not connection.features.supports_transactions')
+    def test_serialized_rollback(
+        self,
+        fixture_with_serialized_rollback: None,
+        fixture_with_db: None,
+    ) -> None:
+        pass
+
 
 class TestDatabaseMarker:
     "Tests for the django_db marker."
@@ -264,6 +332,19 @@ def test_all_databases(self, request) -> None:
         SecondItem.objects.count()
         SecondItem.objects.create(name="spam")
 
+    @pytest.mark.django_db
+    def test_serialized_rollback_disabled(self, request):
+        marker = request.node.get_closest_marker("django_db")
+        assert not marker.kwargs
+
+    # The test works when transactions are not supported, but it interacts
+    # badly with other tests.
+    @pytest.mark.skipif('not connection.features.supports_transactions')
+    @pytest.mark.django_db(serialized_rollback=True)
+    def test_serialized_rollback_enabled(self, request):
+        marker = request.node.get_closest_marker("django_db")
+        assert marker.kwargs["serialized_rollback"]
+
 
 def test_unittest_interaction(django_testdir) -> None:
     "Test that (non-Django) unittests cannot access the DB."
diff --git a/tests/test_db_setup.py b/tests/test_db_setup.py
