Support reset sequences

This adds support for using the reset_sequences feature of Django's
TransactionTestCase.

It will try to reset all automatic increment values before test
execution, if the database supports it. This is useful for when you have
tests that rely on such values, like ids or other primary keys.

Author: cb109

docs/helpers.rst

@@ -13,10 +13,10 @@ 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`` - request database access
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. :py:function:: pytest.mark.django_db:
+.. :py:function:: pytest.mark.django_db([transaction=False, reset_sequences=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
@@ -25,9 +25,9 @@ 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 ``reset_sequences_db`` fixtures.  Otherwise the
+test will fail when trying to access the database.
 
 :type transaction: bool
 :param transaction:
@@ -38,14 +38,23 @@ when trying to access the database.
  uses. When ``transaction=True``, the behavior will be the same as
  `django.test.TransactionTestCase`_
 
+
+:type reset_sequences: bool
+:param reset_sequences:
+ The ``reset_sequences`` argument will ask to reset auto increment sequence
+ values (e.g. primary keys) before running the test.  Defaults to
+ ``False``.  Must be used together with ``transaction=True`` to have an
+ effect.  Please be aware that not all databases support this feature.
+ For details see `django.test.TransactionTestCase.reset_sequences`_
+
 .. 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 ``reset_sequences_db`` fixture.  See below
+  for a description of them.
 
 .. note:: Automatic usage with ``django.test.TestCase``.
 
@@ -54,6 +63,7 @@ when trying to access the database.
  Test classes that subclass Python's ``unittest.TestCase`` need to have the
  marker applied in order to access the database.
 
+.. _django.test.TransactionTestCase.reset_sequences: https://docs.djangoproject.com/en/dev/topics/testing/advanced/#django.test.TransactionTestCase.reset_sequences
 .. _django.test.TestCase: https://docs.djangoproject.com/en/dev/topics/testing/overview/#testcase
 .. _django.test.TransactionTestCase: https://docs.djangoproject.com/en/dev/topics/testing/overview/#transactiontestcase
 
@@ -215,8 +225,17 @@ mark to signal it needs the database.
 
 This fixture can be used to request access to the database including
 transaction support.  This is only required for fixtures which need
-database access themselves.  A test function would normally use the
-``pytest.mark.django_db`` mark to signal it needs the database.
+database access themselves.  A test function should normally use the
+``pytest.mark.django_db``  mark with ``transaction=True``.
+
+``reset_sequences_db``
+~~~~~~~~~~~~~~~~~~~~~~
+
+This fixture provides the same transactional database access as
+``transactional_db``, with additional support for reset of auto increment
+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``.
 
 ``live_server``
 ~~~~~~~~~~~~~~~
@@ -227,6 +246,18 @@ or by requesting it's string value: ``unicode(live_server)``.  You can
 also directly concatenate a string to form a URL: ``live_server +
 '/foo``.
 
+.. note:: Combining database access fixtures.
+
+  When using multiple database fixtures together, only one of them is
+  used.  Their order of precedence is as follows (the last one wins):
+
+  * ``db``
+  * ``transactional_db``
+  * ``reset_sequences_db``
+
+  In addition, using ``live_server`` will also trigger transactional
+  database access, if not specified.
+
 ``settings``
 ~~~~~~~~~~~~
 

pytest_django/fixtures.py

@@ -14,8 +14,8 @@
 
 from .lazy_django import 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', 'django_db_reset_sequences',
+           'admin_user', 'django_user_model', 'django_username_field',
            'client', 'admin_client', 'rf', 'settings', 'live_server',
            '_live_server_helper', 'django_assert_num_queries']
 
@@ -107,7 +107,8 @@ def teardown_database():
         request.addfinalizer(teardown_database)
 
 
-def _django_db_fixture_helper(transactional, request, django_db_blocker):
+def _django_db_fixture_helper(request, django_db_blocker,
+                              transactional=False, reset_sequences=False):
     if is_django_unittest(request):
         return
 
@@ -120,6 +121,11 @@ def _django_db_fixture_helper(transactional, request, django_db_blocker):
 
     if transactional:
         from django.test import TransactionTestCase as django_case
+
+        if reset_sequences:
+            class ResetSequenceTestCase(django_case):
+                reset_sequences = True
+            django_case = ResetSequenceTestCase
     else:
         from django.test import TestCase as django_case
 
@@ -139,7 +145,7 @@ def _disable_native_migrations():
 
 @pytest.fixture(scope='function')
 def db(request, django_db_setup, django_db_blocker):
-    """Require a django test database
+    """Require a django test database.
 
     This database will be setup with the default fixtures and will have
     the transaction management disabled. At the end of the test the outer
@@ -148,30 +154,54 @@ def db(request, django_db_setup, django_db_blocker):
     This is more limited than the ``transactional_db`` resource but
     faster.
 
-    If both this and ``transactional_db`` are requested then the
-    database setup will behave as only ``transactional_db`` was
-    requested.
+    If multiple database fixtures are requested, they take precedence
+    over each other in the following order (the last one wins): ``db``,
+    ``transactional_db``, ``django_db_reset_sequences``.
     """
+    if 'django_db_reset_sequences' in request.funcargnames:
+        request.getfixturevalue('django_db_reset_sequences')
     if 'transactional_db' in request.funcargnames \
             or 'live_server' in request.funcargnames:
         request.getfixturevalue('transactional_db')
     else:
-        _django_db_fixture_helper(False, request, django_db_blocker)
+        _django_db_fixture_helper(request, django_db_blocker, transactional=False)
 
 
 @pytest.fixture(scope='function')
 def transactional_db(request, django_db_setup, django_db_blocker):
-    """Require a django test database with transaction support
+    """Require a django test database with transaction support.
 
     This will re-initialise the django database for each test and is
     thus slower than the normal ``db`` fixture.
 
     If you want to use the database with transactions you must request
-    this resource.  If both this and ``db`` are requested then the
-    database setup will behave as only ``transactional_db`` was
-    requested.
+    this resource.
+
+    If multiple database fixtures are requested, they take precedence
+    over each other in the following order (the last one wins): ``db``,
+    ``transactional_db``, ``django_db_reset_sequences``.
+    """
+    if 'django_db_reset_sequences' in request.funcargnames:
+        request.getfuncargvalue('django_db_reset_sequences')
+    _django_db_fixture_helper(request, django_db_blocker,
+                              transactional=True)
+
+
+@pytest.fixture(scope='function')
+def django_db_reset_sequences(request, django_db_setup, django_db_blocker):
+    """Require a transactional test database with sequence reset support.
+
+    This behaves like the ``transactional_db`` fixture, with the addition
+    of enforcing a reset of all auto increment sequences.  If the enquiring
+    test relies on such values (e.g. ids as primary keys), you should
+    request this resource to ensure they are consistent across tests.
+
+    If multiple database fixtures are requested, they take precedence
+    over each other in the following order (the last one wins): ``db``,
+    ``transactional_db``, ``django_db_reset_sequences``.
     """
-    _django_db_fixture_helper(True, request, django_db_blocker)
+    _django_db_fixture_helper(request, django_db_blocker,
+                              transactional=True, reset_sequences=True)
 
 
 @pytest.fixture()

pytest_django/plugin.py

@@ -30,6 +30,7 @@
 from .fixtures import django_user_model  # noqa
 from .fixtures import django_username_field  # noqa
 from .fixtures import live_server  # noqa
+from .fixtures import django_db_reset_sequences  # noqa
 from .fixtures import rf  # noqa
 from .fixtures import settings  # noqa
 from .fixtures import transactional_db  # noqa
@@ -379,13 +380,15 @@ 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
+    ``django_db_reset_sequences`` fixtures as required by the django_db marker.
     """
     marker = request.node.get_closest_marker('django_db')
     if marker:
-        transaction = validate_django_db(marker)
-        if transaction:
+        transaction, reset_sequences = validate_django_db(marker)
+        if reset_sequences:
+            request.getfixturevalue('django_db_reset_sequences')
+        elif transaction:
             request.getfixturevalue('transactional_db')
         else:
             request.getfixturevalue('db')
@@ -667,11 +670,14 @@ 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
+    ``reset_sequences`` attributes on the marker which will have the
+    correct values.
+
+    A sequence reset is only allowed when combined with a transaction.
     """
-    def apifun(transaction=False):
-        return transaction
+    def apifun(transaction=False, reset_sequences=False):
+        return transaction, reset_sequences
     return apifun(*marker.args, **marker.kwargs)
 
 

tests/test_database.py

@@ -7,6 +7,12 @@
 from pytest_django_test.app.models import Item
 
 
+def db_supports_reset_sequences():
+    """Return if the current db engine supports `reset_sequences`."""
+    return (connection.features.supports_transactions and
+            connection.features.supports_sequence_reset)
+
+
 def test_noaccess():
     with pytest.raises(pytest.fail.Exception):
         Item.objects.create(name='spam')
@@ -27,20 +33,34 @@ def test_noaccess_fixture(noaccess):
     pass
 
 
-class TestDatabaseFixtures:
-    """Tests for the db and transactional_db fixtures"""
+@pytest.fixture
+def non_zero_sequences_counter(db):
+    """Ensure that the db's internal sequence counter is > 1.
+
+    This is used to test the `reset_sequences` feature.
+    """
+    item_1 = Item.objects.create(name='item_1')
+    item_2 = Item.objects.create(name='item_2')
+    item_1.delete()
+    item_2.delete()
+
 
-    @pytest.fixture(params=['db', 'transactional_db'])
-    def both_dbs(self, request):
-        if request.param == 'transactional_db':
-            return request.getfixturevalue('transactional_db')
+class TestDatabaseFixtures:
+    """Tests for the different database fixtures."""
+
+    @pytest.fixture(params=['db', 'transactional_db', 'django_db_reset_sequences'])
+    def all_dbs(self, request):
+        if request.param == 'django_db_reset_sequences':
+            return request.getfuncargvalue('django_db_reset_sequences')
+        elif request.param == 'transactional_db':
+            return request.getfuncargvalue('transactional_db')
         elif request.param == 'db':
             return request.getfixturevalue('db')
 
-    def test_access(self, both_dbs):
+    def test_access(self, all_dbs):
         Item.objects.create(name='spam')
 
-    def test_clean_db(self, both_dbs):
+    def test_clean_db(self, all_dbs):
         # Relies on the order: test_access created an object
         assert Item.objects.count() == 0
 
@@ -56,8 +76,39 @@ def test_transactions_enabled(self, transactional_db):
 
         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')
+
+        assert not connection.in_atomic_block
+
+    def test_django_db_reset_sequences_fixture(
+            self, db, django_testdir, non_zero_sequences_counter):
+
+        if not db_supports_reset_sequences():
+            pytest.skip('transactions and reset_sequences must be supported '
+                        'by the database to run this test')
+
+        # The test runs on a database that already contains objects, so its
+        # id counter is > 1. We check for the ids of newly created objects.
+        django_testdir.create_test_module('''
+            import pytest
+            from .app.models import Item
+
+            def test_django_db_reset_sequences_requested(
+                    django_db_reset_sequences):
+                item = Item.objects.create(name='new_item')
+                assert item.id == 1
+        ''')
+
+        result = django_testdir.runpytest_subprocess('-v', '--reuse-db')
+        result.stdout.fnmatch_lines([
+            "*test_django_db_reset_sequences_requested PASSED*",
+        ])
+
     @pytest.fixture
-    def mydb(self, both_dbs):
+    def mydb(self, all_dbs):
         # This fixture must be able to access the database
         Item.objects.create(name='spam')
 
@@ -69,13 +120,13 @@ def test_mydb(self, mydb):
         item = Item.objects.get(name='spam')
         assert item
 
-    def test_fixture_clean(self, both_dbs):
+    def test_fixture_clean(self, all_dbs):
         # Relies on the order: test_mydb created an object
         # See https://github.com/pytest-dev/pytest-django/issues/17
         assert Item.objects.count() == 0
 
     @pytest.fixture
-    def fin(self, request, both_dbs):
+    def fin(self, request, all_dbs):
         # This finalizer must be able to access the database
         request.addfinalizer(lambda: Item.objects.create(name='spam'))
 
@@ -139,6 +190,18 @@ def test_transactions_enabled(self):
 
         assert not connection.in_atomic_block
 
+    @pytest.mark.django_db
+    def test_reset_sequences_disabled(self, request):
+        marker = request.keywords['django_db']
+
+        assert not marker.kwargs
+
+    @pytest.mark.django_db(reset_sequences=True)
+    def test_reset_sequences_enabled(self, request):
+        marker = request.keywords['django_db']
+
+        assert marker.kwargs['reset_sequences']
+
 
 def test_unittest_interaction(django_testdir):
     "Test that (non-Django) unittests cannot access the DB."