Skip to content

Commit 4d982e2

Browse files
committed
Added support for multi_db
Created a `django_multi_db` fixture and added a `multi_db` kwarg to the `django_db` marker.
1 parent 5c91295 commit 4d982e2

File tree

5 files changed

+111
-9
lines changed

5 files changed

+111
-9
lines changed

docs/helpers.rst

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ on what marks are and for notes on using_ them.
1616
``pytest.mark.django_db`` - request database access
1717
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1818

19-
.. :py:function:: pytest.mark.django_db([transaction=False, reset_sequences=False]):
19+
.. :py:function:: pytest.mark.django_db([transaction=False, reset_sequences=False, multi_db=False]):
2020
2121
This is used to mark a test function as requiring the database. It
2222
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
2626

2727
In order for a test to have access to the database it must either
2828
be marked using the ``django_db`` mark or request one of the ``db``,
29-
``transactional_db`` or ``django_db_reset_sequences`` fixtures. Otherwise the
30-
test will fail when trying to access the database.
29+
``transactional_db``, ``django_db_reset_sequences`` or ``multi_db`` fixtures.
30+
Otherwise the test will fail when trying to access the database.
3131

3232
:type transaction: bool
3333
:param transaction:
@@ -47,6 +47,13 @@ test will fail when trying to access the database.
4747
effect. Please be aware that not all databases support this feature.
4848
For details see :py:attr:`django.test.TransactionTestCase.reset_sequences`.
4949

50+
51+
:type multi_db: bool
52+
:param multi_db:
53+
The ``multi_db`` argument will allow to test using multiple databases.
54+
This behaves the same way the ``multi_db`` parameter of `django.test.TestCase`_
55+
does.
56+
5057
.. note::
5158

5259
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
242249
fixtures which need database access themselves. A test function should
243250
normally use the ``pytest.mark.django_db`` mark with ``transaction=True`` and ``reset_sequences=True``.
244251

252+
``django_multi_db``
253+
~~~~~~~~~~~~~~~~~~~
254+
255+
.. fixture:: django_multi_db
256+
257+
This fixtures lets you test against multiple databases. When this fixture
258+
is used, the test behaves as a django TestCase class marked with ``multi_db=True``
259+
does. A test function should normally use the ``pytest.mark.django_db`` mark
260+
with ``multi_db=True``.
261+
245262
``live_server``
246263
~~~~~~~~~~~~~~~
247264

pytest_django/fixtures.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
"db",
1919
"transactional_db",
2020
"django_db_reset_sequences",
21+
"django_multi_db",
2122
"admin_user",
2223
"django_user_model",
2324
"django_username_field",
@@ -141,6 +142,13 @@ class ResetSequenceTestCase(django_case):
141142
else:
142143
from django.test import TestCase as django_case
143144

145+
# We check if the multi_db marker has been used
146+
marker = request.node.get_closest_marker('django_db')
147+
multi_db = marker.kwargs.get('multi_db', False) if marker else False
148+
# We check if django_multi_db fixture has been used
149+
multi_db = multi_db or "django_multi_db" in request.fixturenames
150+
django_case.multi_db = multi_db
151+
144152
test_case = django_case(methodName="__init__")
145153
test_case._pre_setup()
146154
request.addfinalizer(test_case._post_teardown)
@@ -219,6 +227,19 @@ def django_db_reset_sequences(request, django_db_setup, django_db_blocker):
219227
)
220228

221229

230+
@pytest.fixture(scope="function")
231+
def django_multi_db(request, django_db_setup, django_db_blocker):
232+
"""Require a django test database
233+
234+
This behaves like the ``db`` fixture, with the addition of marking
235+
the test as multi_db for django test case purposes. Using this fixture
236+
is equivalent to marking your TestCase class as ``multi_db = True``.
237+
238+
You can use this fixture in tandem with other fixtures.
239+
"""
240+
request.getfixturevalue("db")
241+
242+
222243
@pytest.fixture()
223244
def client():
224245
"""A Django test client instance."""

pytest_django/plugin.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from .fixtures import rf # noqa
3636
from .fixtures import settings # noqa
3737
from .fixtures import transactional_db # noqa
38+
from .fixtures import django_multi_db # noqa
3839

3940
from .lazy_django import django_settings_is_configured, skip_if_no_django
4041

@@ -495,12 +496,14 @@ def django_db_blocker():
495496
def _django_db_marker(request):
496497
"""Implement the django_db marker, internal to pytest-django.
497498
498-
This will dynamically request the ``db``, ``transactional_db`` or
499-
``django_db_reset_sequences`` fixtures as required by the django_db marker.
499+
This will dynamically request the ``db``, ``transactional_db``,
500+
``django_db_reset_sequences`` or ``multi_db`` fixtures as
501+
required by the django_db marker.
500502
"""
501503
marker = request.node.get_closest_marker("django_db")
502504
if marker:
503-
transaction, reset_sequences = validate_django_db(marker)
505+
transaction, reset_sequences, multi_db = validate_django_db(marker)
506+
# multi_db is handled in `_django_db_fixture_helper`
504507
if reset_sequences:
505508
request.getfixturevalue("django_db_reset_sequences")
506509
elif transaction:
@@ -804,8 +807,8 @@ def validate_django_db(marker):
804807
A sequence reset is only allowed when combined with a transaction.
805808
"""
806809

807-
def apifun(transaction=False, reset_sequences=False):
808-
return transaction, reset_sequences
810+
def apifun(transaction=False, reset_sequences=False, multi_db=False):
811+
return transaction, reset_sequences, multi_db
809812

810813
return apifun(*marker.args, **marker.kwargs)
811814

pytest_django_test/settings_sqlite.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,5 +4,8 @@
44
"default": {
55
"ENGINE": "django.db.backends.sqlite3",
66
"NAME": "/should_not_be_accessed",
7-
}
7+
},
88
}
9+
10+
DATABASES["replica"] = DATABASES["default"].copy()
11+
DATABASES["replica"]["NAME"] += '_replica'

tests/test_database.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,12 @@ def test_transactions_enabled(self, transactional_db):
7878

7979
assert not connection.in_atomic_block
8080

81+
def test_transactions_enabled_multi_db(self, transactional_db, django_multi_db):
82+
if not connections_support_transactions():
83+
pytest.skip("transactions required for this test")
84+
85+
assert not connection.in_atomic_block
86+
8187
def test_transactions_enabled_via_reset_seq(self, django_db_reset_sequences):
8288
if not connections_support_transactions():
8389
pytest.skip("transactions required for this test")
@@ -140,6 +146,22 @@ def test_fin(self, fin):
140146
# Check finalizer has db access (teardown will fail if not)
141147
pass
142148

149+
def test_multi_db_access(self, all_dbs, django_multi_db):
150+
Item.objects.using('replica').create(name="spam")
151+
152+
def test_multi_db_clean(self, all_dbs, django_multi_db):
153+
# Relies on the order: test_multi_db_access created an object
154+
assert Item.objects.using('replica').count() == 0
155+
156+
def test_no_multi_db_access(self, all_dbs):
157+
# Even without marker we can write to replica
158+
# but items won't be cleaned, see `test_no_multi_db_no_clean`
159+
Item.objects.using('replica').create(name="spam")
160+
161+
def test_no_multi_db_no_clean(self, all_dbs):
162+
# Relies on the order: test_no_multi_db_access created objects
163+
assert Item.objects.using('replica').count() > 0
164+
143165

144166
class TestDatabaseFixturesAllOrder:
145167
@pytest.fixture
@@ -205,6 +227,13 @@ def test_transactions_enabled(self):
205227

206228
assert not connection.in_atomic_block
207229

230+
@pytest.mark.django_db(transaction=True, multi_db=True)
231+
def test_transactions_enabled_multi_db(self):
232+
if not connections_support_transactions():
233+
pytest.skip("transactions required for this test")
234+
235+
assert not connection.in_atomic_block
236+
208237
@pytest.mark.django_db
209238
def test_reset_sequences_disabled(self, request):
210239
marker = request.node.get_closest_marker("django_db")
@@ -215,6 +244,35 @@ def test_reset_sequences_enabled(self, request):
215244
marker = request.node.get_closest_marker("django_db")
216245
assert marker.kwargs["reset_sequences"]
217246

247+
@pytest.mark.django_db(multi_db=True)
248+
def test_access_multi_db(self):
249+
Item.objects.using('replica').create(name="spam")
250+
251+
@pytest.mark.django_db(multi_db=True)
252+
def test_clean_multi_db(self):
253+
# Relies on the order: test_access_multi_db created an object.
254+
assert Item.objects.using('replica').count() == 0
255+
256+
@pytest.mark.django_db(transaction=True, multi_db=True)
257+
def test_transaction_access_multi_db(self):
258+
Item.objects.using('replica').create(name="spam")
259+
260+
@pytest.mark.django_db(transaction=True, multi_db=True)
261+
def test_transaction_clean_multi_db(self):
262+
# Relies on the order: test_transaction_access_multi_db created an object.
263+
assert Item.objects.using('replica').count() == 0
264+
265+
@pytest.mark.django_db
266+
def test_no_multi_db_access(self):
267+
# Even without marker we can write to replica
268+
# but items won't be cleaned, see `test_no_multi_db_no_clean`
269+
Item.objects.using('replica').create(name="spam")
270+
271+
@pytest.mark.django_db
272+
def test_no_multi_db_clean(self):
273+
# Relies on the order: test_no_multi_db_access created objects
274+
assert Item.objects.using('replica').count() > 0
275+
218276

219277
def test_unittest_interaction(django_testdir):
220278
"Test that (non-Django) unittests cannot access the DB."

0 commit comments

Comments
 (0)