Added support for multi_db

docs/helpers.rst

@@ -16,7 +16,7 @@ on what marks are and for notes on using_ them.
 ``pytest.mark.django_db`` - request database access
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. :py:function:: pytest.mark.django_db([transaction=False, reset_sequences=False]):
+.. :py:function:: pytest.mark.django_db([transaction=False, reset_sequences=False, multi_db=False]):
 
 This is used to mark a test function as requiring the database. It
 will ensure the database is set up correctly for the test. Each test
@@ -26,8 +26,8 @@ of the test. This behavior is the same as Django's standard
 
 In order for a test to have access to the database it must either
 be marked using the ``django_db`` mark or request one of the ``db``,
-``transactional_db`` or ``django_db_reset_sequences`` fixtures.  Otherwise the
-test will fail when trying to access the database.
+``transactional_db``, ``django_db_reset_sequences`` or ``multi_db`` fixtures.
+Otherwise the test will fail when trying to access the database.
 
 :type transaction: bool
 :param transaction:
@@ -47,6 +47,13 @@ test will fail when trying to access the database.
  effect.  Please be aware that not all databases support this feature.
  For details see :py:attr:`django.test.TransactionTestCase.reset_sequences`.
 
+
+:type multi_db: bool
+:param multi_db:
+ The ``multi_db`` argument will allow to test using multiple databases.
+ This behaves the same way the ``multi_db`` parameter of `django.test.TestCase`_
+ does.
+
 .. note::
 
   If you want access to the Django database *inside a fixture*
@@ -242,6 +249,16 @@ sequences (if your database supports it). This is only required for
 fixtures which need database access themselves. A test function should
 normally use the ``pytest.mark.django_db`` mark with ``transaction=True`` and ``reset_sequences=True``.
 
+``django_multi_db``
+~~~~~~~~~~~~~~~~~~~
+
+.. fixture:: django_multi_db
+
+This fixtures lets you test against multiple databases. When this fixture
+is used, the test behaves as a django TestCase class marked with ``multi_db=True``
+does. A test function should normally use the ``pytest.mark.django_db`` mark
+with ``multi_db=True``.
+
 ``live_server``
 ~~~~~~~~~~~~~~~
 

pytest_django/fixtures.py

@@ -18,6 +18,7 @@
     "db",
     "transactional_db",
     "django_db_reset_sequences",
+    "django_multi_db",
     "admin_user",
     "django_user_model",
     "django_username_field",
@@ -141,6 +142,13 @@ class ResetSequenceTestCase(django_case):
     else:
         from django.test import TestCase as django_case
 
+    # We check if the multi_db marker has been used
+    marker = request.node.get_closest_marker('django_db')
+    multi_db = marker.kwargs.get('multi_db', False) if marker else False
+    # We check if django_multi_db fixture has been used
+    multi_db = multi_db or "django_multi_db" in request.fixturenames
+    django_case.multi_db = multi_db
+
     test_case = django_case(methodName="__init__")
     test_case._pre_setup()
     request.addfinalizer(test_case._post_teardown)
@@ -219,6 +227,19 @@ def django_db_reset_sequences(request, django_db_setup, django_db_blocker):
     )
 
 
+@pytest.fixture(scope="function")
+def django_multi_db(request, django_db_setup, django_db_blocker):
+    """Require a django test database
+
+    This behaves like the ``db`` fixture, with the addition of marking
+    the test as multi_db for django test case purposes. Using this fixture
+    is equivalent to marking your TestCase class as ``multi_db = True``.
+
+    You can use this fixture in tandem with other fixtures.
+    """
+    request.getfixturevalue("db")
+
+
 @pytest.fixture()
 def client():
     """A Django test client instance."""

pytest_django/plugin.py

@@ -35,6 +35,7 @@
 from .fixtures import rf  # noqa
 from .fixtures import settings  # noqa
 from .fixtures import transactional_db  # noqa
+from .fixtures import django_multi_db  # noqa
 
 from .lazy_django import django_settings_is_configured, skip_if_no_django
 
@@ -495,12 +496,14 @@ def django_db_blocker():
 def _django_db_marker(request):
     """Implement the django_db marker, internal to pytest-django.
 
-    This will dynamically request the ``db``, ``transactional_db`` or
-    ``django_db_reset_sequences`` fixtures as required by the django_db marker.
+    This will dynamically request the ``db``, ``transactional_db``,
+    ``django_db_reset_sequences`` or ``multi_db`` fixtures as
+    required by the django_db marker.
     """
     marker = request.node.get_closest_marker("django_db")
     if marker:
-        transaction, reset_sequences = validate_django_db(marker)
+        transaction, reset_sequences, multi_db = validate_django_db(marker)
+        # multi_db is handled in `_django_db_fixture_helper`
         if reset_sequences:
             request.getfixturevalue("django_db_reset_sequences")
         elif transaction:
@@ -804,8 +807,8 @@ def validate_django_db(marker):
     A sequence reset is only allowed when combined with a transaction.
     """
 
-    def apifun(transaction=False, reset_sequences=False):
-        return transaction, reset_sequences
+    def apifun(transaction=False, reset_sequences=False, multi_db=False):
+        return transaction, reset_sequences, multi_db
 
     return apifun(*marker.args, **marker.kwargs)
 

pytest_django_test/settings_sqlite.py

@@ -4,5 +4,8 @@
     "default": {
         "ENGINE": "django.db.backends.sqlite3",
         "NAME": "/should_not_be_accessed",
-    }
+    },
 }
+
+DATABASES["replica"] = DATABASES["default"].copy()
+DATABASES["replica"]["NAME"] += '_replica'

tests/test_database.py

@@ -78,6 +78,12 @@ def test_transactions_enabled(self, transactional_db):
 
         assert not connection.in_atomic_block
 
+    def test_transactions_enabled_multi_db(self, transactional_db, django_multi_db):
+        if not connections_support_transactions():
+            pytest.skip("transactions required for this test")
+
+        assert not connection.in_atomic_block
+
     def test_transactions_enabled_via_reset_seq(self, django_db_reset_sequences):
         if not connections_support_transactions():
             pytest.skip("transactions required for this test")
@@ -140,6 +146,22 @@ def test_fin(self, fin):
         # Check finalizer has db access (teardown will fail if not)
         pass
 
+    def test_multi_db_access(self, all_dbs, django_multi_db):
+        Item.objects.using('replica').create(name="spam")
+
+    def test_multi_db_clean(self, all_dbs, django_multi_db):
+        # Relies on the order: test_multi_db_access created an object
+        assert Item.objects.using('replica').count() == 0
+
+    def test_no_multi_db_access(self, all_dbs):
+        # Even without marker we can write to replica
+        # but items won't be cleaned, see `test_no_multi_db_no_clean`
+        Item.objects.using('replica').create(name="spam")
+
+    def test_no_multi_db_no_clean(self, all_dbs):
+        # Relies on the order: test_no_multi_db_access created objects
+        assert Item.objects.using('replica').count() > 0
+
 
 class TestDatabaseFixturesAllOrder:
     @pytest.fixture
@@ -205,6 +227,13 @@ def test_transactions_enabled(self):
 
         assert not connection.in_atomic_block
 
+    @pytest.mark.django_db(transaction=True, multi_db=True)
+    def test_transactions_enabled_multi_db(self):
+        if not connections_support_transactions():
+            pytest.skip("transactions required for this test")
+
+        assert not connection.in_atomic_block
+
     @pytest.mark.django_db
     def test_reset_sequences_disabled(self, request):
         marker = request.node.get_closest_marker("django_db")
@@ -215,6 +244,35 @@ def test_reset_sequences_enabled(self, request):
         marker = request.node.get_closest_marker("django_db")
         assert marker.kwargs["reset_sequences"]
 
+    @pytest.mark.django_db(multi_db=True)
+    def test_access_multi_db(self):
+        Item.objects.using('replica').create(name="spam")
+
+    @pytest.mark.django_db(multi_db=True)
+    def test_clean_multi_db(self):
+        # Relies on the order: test_access_multi_db created an object.
+        assert Item.objects.using('replica').count() == 0
+
+    @pytest.mark.django_db(transaction=True, multi_db=True)
+    def test_transaction_access_multi_db(self):
+        Item.objects.using('replica').create(name="spam")
+
+    @pytest.mark.django_db(transaction=True, multi_db=True)
+    def test_transaction_clean_multi_db(self):
+        # Relies on the order: test_transaction_access_multi_db created an object.
+        assert Item.objects.using('replica').count() == 0
+
+    @pytest.mark.django_db
+    def test_no_multi_db_access(self):
+        # Even without marker we can write to replica
+        # but items won't be cleaned, see `test_no_multi_db_no_clean`
+        Item.objects.using('replica').create(name="spam")
+
+    @pytest.mark.django_db
+    def test_no_multi_db_clean(self):
+        # Relies on the order: test_no_multi_db_access created objects
+        assert Item.objects.using('replica').count() > 0
+
 
 def test_unittest_interaction(django_testdir):
     "Test that (non-Django) unittests cannot access the DB."