Add support for multi_db

docs/helpers.rst

@@ -13,8 +13,8 @@ on what marks are and for notes on using_ them.
 .. _using: https://pytest.org/en/latest/example/markers.html#marking-whole-classes-or-modules
 
 
-``pytest.mark.django_db(transaction=False)`` - request database access
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+``pytest.mark.django_db(transaction=False, multi_db=False)`` - request database access
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. :py:function:: pytest.mark.django_db:
 
@@ -25,27 +25,35 @@ of the test. This behavior is the same as Django's standard
 `django.test.TestCase`_ class.
 
 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``
-or ``transactional_db`` fixtures.  Otherwise the test will fail
-when trying to access the database.
+be marked using the ``django_db`` mark or request one of the ``db``,
+``transactional_db`` or ``use_multi_db`` fixtures.  Otherwise the test will
+fail when trying to access the database.
 
 :type transaction: bool
 :param transaction:
- The ``transaction`` argument will allow the test to use real transactions.
- With ``transaction=False`` (the default when not specified), transaction
- operations are noops during the test. This is the same behavior that
- `django.test.TestCase`_
- uses. When ``transaction=True``, the behavior will be the same as
- `django.test.TransactionTestCase`_
+  The ``transaction`` argument will allow the test to use real transactions.
+  With ``transaction=False`` (the default when not specified), transaction
+  operations are noops during the test. This is the same behavior that
+  `django.test.TestCase`_
+  uses. When ``transaction=True``, the behavior will be the same as
+  `django.test.TransactionTestCase`_
+
+:type multi_db: bool
+:param multi_db:
+  The ``multi_db`` allows the test to flush *all* databases before running.
+  With ``multi_db=False`` (the default when not specified), only the *default*
+  database will be flushed. This is the same behavior that `django.test.TestCase`_
+  uses. When ``multi_db=True``, the behavior will be the same as a
+  `django.test.TestCase`_ with the class attribute ``multi_db = True``
 
 .. note::
 
   If you want access to the Django database *inside a fixture*
   this marker will not help even if the function requesting your
   fixture has this marker applied.  To access the database in a
-  fixture, the fixture itself will have to request the ``db`` or
-  ``transactional_db`` fixture.  See below for a description of
-  them.
+  fixture, the fixture itself will have to request the ``db``,
+  ``transactional_db`` or ``use_multi_db`` fixture.  See below for a
+  description of them.
 
 .. note:: Automatic usage with ``django.test.TestCase``.
 
@@ -193,6 +201,15 @@ transaction support.  This is only required for fixtures which need
 database access themselves.  A test function would normally use the
 :py:func:`~pytest.mark.django_db` mark to signal it needs the database.
 
+``use_multi_db``
+~~~~~~~~~~~~
+
+This fixture can be used to request access to the database including
+multiple database support.  This is only required for fixtures which
+need databases access themselves.  A test function would normally use the
+:py:func:`~pytest.mark.django_db` mark to signal it needs the database.
+
+
 ``live_server``
 ~~~~~~~~~~~~~~~
 

pytest_django/fixtures.py

@@ -15,8 +15,8 @@
 
 from .lazy_django import get_django_version, skip_if_no_django
 
-__all__ = ['django_db_setup', 'db', 'transactional_db', 'admin_user',
-           'django_user_model', 'django_username_field',
+__all__ = ['django_db_setup', 'db', 'transactional_db', 'use_multi_db',
+           'admin_user', 'django_user_model', 'django_username_field',
            'client', 'admin_client', 'rf', 'settings', 'live_server',
            '_live_server_helper', 'django_assert_num_queries']
 
@@ -110,7 +110,8 @@ def teardown_database():
         request.addfinalizer(teardown_database)
 
 
-def _django_db_fixture_helper(transactional, request, django_db_blocker):
+def _django_db_fixture_helper(transactional, request,
+                              django_db_blocker, multi_db=False):
     if is_django_unittest(request):
         return
 
@@ -126,6 +127,7 @@ def _django_db_fixture_helper(transactional, request, django_db_blocker):
     else:
         from django.test import TestCase as django_case
 
+    django_case.multi_db = multi_db
     test_case = django_case(methodName='__init__')
     test_case._pre_setup()
     request.addfinalizer(test_case._post_teardown)
@@ -177,6 +179,25 @@ def transactional_db(request, django_db_setup, django_db_blocker):
     _django_db_fixture_helper(True, request, django_db_blocker)
 
 
+@pytest.fixture(scope='function')
+def use_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``.
+
+    If both this and ``transactional_db`` are requested then the
+    database setup will behave as only ``transactional_db`` was
+    requested.
+    """
+    if 'transactional_db' in request.funcargnames \
+            or 'live_server' in request.funcargnames:
+        request.getfuncargvalue('transactional_db')
+    else:
+        _django_db_fixture_helper(False, request, django_db_blocker, multi_db=True)
+
+
 @pytest.fixture()
 def client():
     """A Django test client instance."""

pytest_django/plugin.py

@@ -32,6 +32,7 @@
 from .fixtures import rf  # noqa
 from .fixtures import settings  # noqa
 from .fixtures import transactional_db  # noqa
+from .fixtures import use_multi_db  # noqa
 from .pytest_compat import getfixturevalue
 
 from .lazy_django import (django_settings_is_configured,
@@ -365,14 +366,16 @@ def django_db_blocker():
 def _django_db_marker(request):
     """Implement the django_db marker, internal to pytest-django.
 
-    This will dynamically request the ``db`` or ``transactional_db``
-    fixtures as required by the django_db marker.
+    This will dynamically request the ``db``, ``transactional_db``
+    or ``use_multi_db`` fixtures as required by the django_db marker.
     """
     marker = request.keywords.get('django_db', None)
     if marker:
         validate_django_db(marker)
         if marker.transaction:
             getfixturevalue(request, 'transactional_db')
+        elif marker.multi_db:
+            getfixturevalue(request, 'use_multi_db')
         else:
             getfixturevalue(request, 'db')
 
@@ -642,11 +645,12 @@ def restore(self):
 def validate_django_db(marker):
     """Validate the django_db marker.
 
-    It checks the signature and creates the `transaction` attribute on
-    the marker which will have the correct value.
+    It checks the signature and creates the `transaction` and `multi_db`
+    attributes on the marker which will have the correct value.
     """
-    def apifun(transaction=False):
+    def apifun(transaction=False, multi_db=False):
         marker.transaction = transaction
+        marker.multi_db = multi_db
     apifun(*marker.args, **marker.kwargs)
 
 

tests/test_database.py

@@ -31,12 +31,14 @@ def test_noaccess_fixture(noaccess):
 class TestDatabaseFixtures:
     """Tests for the db and transactional_db fixtures"""
 
-    @pytest.fixture(params=['db', 'transactional_db'])
+    @pytest.fixture(params=['db', 'transactional_db', 'use_multi_db'])
     def both_dbs(self, request):
         if request.param == 'transactional_db':
             return getfixturevalue(request, 'transactional_db')
         elif request.param == 'db':
             return getfixturevalue(request, 'db')
+        elif request.param == 'use_multi_db':
+            return getfixturevalue(request, 'use_multi_db')
 
     def test_access(self, both_dbs):
         Item.objects.create(name='spam')
@@ -57,6 +59,16 @@ def test_transactions_enabled(self, transactional_db):
 
         assert not connection.in_atomic_block
 
+    def test_multi_db_enabled(self, use_multi_db):
+        assert self._is_multi_db()
+
+    def test_multi_db_transactions_enabled(self, transactional_db, use_multi_db):
+        assert self._is_multi_db()
+
+    def _is_multi_db(self):
+        from django.test import TestCase
+        return TestCase.multi_db
+
     @pytest.fixture
     def mydb(self, both_dbs):
         # This fixture must be able to access the database
@@ -140,6 +152,22 @@ def test_transactions_enabled(self):
 
         assert not connection.in_atomic_block
 
+    @pytest.mark.django_db(multi_db=False)
+    def test_multi_db_disabled_explicit(self):
+        assert not self._is_multi_db()
+
+    @pytest.mark.django_db(multi_db=True)
+    def test_multi_db_enabled(self):
+        assert self._is_multi_db()
+
+    @pytest.mark.django_db(multi_db=True, transaction=True)
+    def test_transactions_enabled_multi_db_enabled(self):
+        assert self._is_multi_db()
+
+    def _is_multi_db(self):
+        from django.test import TestCase
+        return TestCase.multi_db
+
 
 def test_unittest_interaction(django_testdir):
     "Test that (non-Django) unittests cannot access the DB."