Skip to content

Commit

Permalink
Use block()/unblock() as method names in django_db_blocker.
Browse files Browse the repository at this point in the history
Also, make both block() and unblock() be context managers and remove the
implicit context manager of django_db_blocker itself.

Fixes #372.
  • Loading branch information
pelme committed Aug 21, 2016
1 parent 64673c4 commit 890a4fb
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 27 deletions.
16 changes: 8 additions & 8 deletions docs/database.rst
Expand Up @@ -252,22 +252,22 @@ access for the specified block::

@pytest.fixture
def myfixture(django_db_blocker):
with django_db_blocker:
with django_db_blocker.unblock():
... # modify something in the database

You can also manage the access manually via these methods:

.. py:method:: django_db_blocker.enable_database_access()
.. py:method:: django_db_blocker.unblock()
Enable database access. Should be followed by a call to
:func:`~django_db_blocker.restore_previous_access`.
:func:`~django_db_blocker.restore`.

.. py:method:: django_db_blocker.disable_database_access()
.. py:method:: django_db_blocker.block()
Disable database access. Should be followed by a call to
:func:`~django_db_blocker.restore_previous_access`.
:func:`~django_db_blocker.restore`.

.. py:function:: django_db_blocker.restore_previous_access()
.. py:function:: django_db_blocker.restore()
Restore the previous state of the database blocking.

Expand Down Expand Up @@ -359,7 +359,7 @@ Put this in ``conftest.py``::

@pytest.fixture(scope='session')
def django_db_setup(django_db_setup, django_db_blocker):
with django_db_blocker:
with django_db_blocker.unblock():
call_command('loaddata', 'your_data_fixture.json')

Use the same database for all xdist processes
Expand Down Expand Up @@ -396,7 +396,7 @@ Put this in ``conftest.py``::

@pytest.fixture(scope='session')
def django_db_setup(django_db_setup, django_db_blocker):
with django_db_blocker:
with django_db_blocker.unblock():
cur = connection.cursor()
cur.execute('ALTER SEQUENCE app_model_id_seq RESTART WITH %s;',
[random.randint(10000, 20000)])
10 changes: 5 additions & 5 deletions pytest_django/fixtures.py
Expand Up @@ -86,18 +86,18 @@ def django_db_setup(
# Django 1.7 compatibility
from .db_reuse import monkey_patch_creation_for_db_reuse

with django_db_blocker:
with django_db_blocker.unblock():
monkey_patch_creation_for_db_reuse()

with django_db_blocker:
with django_db_blocker.unblock():
db_cfg = setup_databases(
verbosity=pytest.config.option.verbose,
interactive=False,
**setup_databases_args
)

def teardown_database():
with django_db_blocker:
with django_db_blocker.unblock():
teardown_databases(
db_cfg,
verbosity=pytest.config.option.verbose,
Expand All @@ -115,8 +115,8 @@ def _django_db_fixture_helper(transactional, request, django_db_blocker):
# Do nothing, we get called with transactional=True, too.
return

django_db_blocker.enable_database_access()
request.addfinalizer(django_db_blocker.restore_previous_access)
django_db_blocker.unblock()
request.addfinalizer(django_db_blocker.restore)

if transactional:
from django.test import TransactionTestCase as django_case
Expand Down
34 changes: 20 additions & 14 deletions pytest_django/plugin.py
Expand Up @@ -146,7 +146,7 @@ def _setup_django():
return

django.setup()
_blocking_manager.disable_database_access()
_blocking_manager.block()


def _get_boolean_value(x, name, default=None):
Expand Down Expand Up @@ -348,8 +348,7 @@ def django_db_blocker():
special database handling.
The object is a context manager and provides the methods
.enable_database_access()/.disable_database_access() and
.restore_database_access() to temporarily enable database access.
.unblock()/.block() and .restore() to temporarily enable database access.
This is an advanced feature that is meant to be used to implement database
fixtures.
Expand Down Expand Up @@ -383,7 +382,7 @@ def _django_setup_unittest(request, django_db_blocker):
request.getfuncargvalue('django_test_environment')
request.getfuncargvalue('django_db_setup')

django_db_blocker.enable_database_access()
django_db_blocker.unblock()

cls = request.node.cls

Expand All @@ -394,7 +393,7 @@ def _django_setup_unittest(request, django_db_blocker):
def teardown():
_restore_class_methods(cls)
cls.tearDownClass()
django_db_blocker.restore_previous_access()
django_db_blocker.restore()

request.addfinalizer(teardown)

Expand Down Expand Up @@ -518,6 +517,17 @@ def _template_string_if_invalid_marker(request):
# ############### Helper Functions ################


class _DatabaseBlockerContextManager(object):
def __init__(self, db_blocker):
self._db_blocker = db_blocker

def __enter__(self):
pass

def __exit__(self, exc_type, exc_value, traceback):
self._db_blocker.restore()


class _DatabaseBlocker(object):
"""Manager for django.db.backends.base.base.BaseDatabaseWrapper.
Expand Down Expand Up @@ -548,25 +558,21 @@ def _blocking_wrapper(*args, **kwargs):
pytest.fail('Database access not allowed, '
'use the "django_db" mark to enable it.')

def enable_database_access(self):
def unblock(self):
"""Enable access to the Django database."""
self._save_active_wrapper()
self._dj_db_wrapper.ensure_connection = self._real_ensure_connection
return _DatabaseBlockerContextManager(self)

def disable_database_access(self):
def block(self):
"""Disable access to the Django database."""
self._save_active_wrapper()
self._dj_db_wrapper.ensure_connection = self._blocking_wrapper
return _DatabaseBlockerContextManager(self)

def restore_previous_access(self):
def restore(self):
self._dj_db_wrapper.ensure_connection = self._history.pop()

def __enter__(self):
self.enable_database_access()

def __exit__(self, exc_type, exc_value, traceback):
self.restore_previous_access()


_blocking_manager = _DatabaseBlocker()

Expand Down
28 changes: 28 additions & 0 deletions tests/test_fixtures.py
Expand Up @@ -304,3 +304,31 @@ class Migration(migrations.Migration):
result = django_testdir.runpytest_subprocess('-s')
result.stdout.fnmatch_lines(['*1 passed*'])
assert result.ret == 0


class Test_django_db_blocker:
@pytest.mark.django_db
def test_block_manually(self, django_db_blocker):
try:
django_db_blocker.block()
with pytest.raises(pytest.fail.Exception):
Item.objects.exists()
finally:
django_db_blocker.restore()

@pytest.mark.django_db
def test_block_with_block(self, django_db_blocker):
with django_db_blocker.block():
with pytest.raises(pytest.fail.Exception):
Item.objects.exists()

def test_unblock_manually(self, django_db_blocker):
try:
django_db_blocker.unblock()
Item.objects.exists()
finally:
django_db_blocker.restore()

def test_unblock_with_block(self, django_db_blocker):
with django_db_blocker.unblock():
Item.objects.exists()

0 comments on commit 890a4fb

Please sign in to comment.