Improve django_db_blocker (#380)

* Move BaseDatabaseWrapper import check to compat
* Use block()/unblock() as method names in django_db_blocker.
* Make both block() and unblock() be context managers and remove the
implicit context manager of django_db_blocker itself.

Fixes #372.

Author: pelme

docs/database.rst

@@ -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.
 
@@ -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
@@ -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)])
@@ -418,7 +418,7 @@ Put this in ``conftest.py``::
 
     @pytest.fixture(scope='session')
     def django_db_setup(django_db_blocker):
-        with django_db_blocker:
+        with django_db_blocker.unblock():
             with connection.cursor() as c:
                 c.executescript('''
                 DROP TABLE IF EXISTS theapp_item;

pytest_django/compat.py

@@ -9,3 +9,9 @@ def teardown_databases(db_cfg, verbosity):
         (_DiscoverRunner(verbosity=verbosity,
                          interactive=False)
          .teardown_databases(db_cfg))
+
+try:
+    from django.db.backends.base.base import BaseDatabaseWrapper  # noqa
+except ImportError:
+    # Django 1.7.
+    from django.db.backends import BaseDatabaseWrapper  # noqa

pytest_django/fixtures.py

@@ -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,
@@ -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

pytest_django/plugin.py

@@ -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):
@@ -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.
@@ -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
 
@@ -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)
 
@@ -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.
 
@@ -530,11 +540,7 @@ def __init__(self):
 
     @property
     def _dj_db_wrapper(self):
-        try:
-            from django.db.backends.base.base import BaseDatabaseWrapper
-        except ImportError:
-            # Django 1.7.
-            from django.db.backends import BaseDatabaseWrapper
+        from .compat import BaseDatabaseWrapper
 
         # The first time the _dj_db_wrapper is accessed, we will save a
         # reference to the real implementation.
@@ -552,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()
 

tests/test_fixtures.py

@@ -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()