Pytest plugin

Author: jpic

dbdiff/fixture.py

@@ -119,7 +119,7 @@ def dump(self, out):
             traceback=True,
             indent=self.indent,
             stdout=out,
-            use_natural_foreign_keys=True,
+            use_natural_foreign_keys=True
         )
 
     def assertNoDiff(self, exclude=None):  # noqa

dbdiff/plugin.py

@@ -0,0 +1,37 @@
+"""Pytest plugin for django-dbdiff.
+
+The marker enables the smarter sequence reset feature previously available in
+the DbdiffTestMixin in pytest, example usage::
+
+    @dbdiff(models=[YourModel])
+    def your_test():
+        assert YourModel.objects.create().pk == 1
+"""
+import pytest
+
+from pytest_django.pytest_compat import getfixturevalue
+
+from .sequence import sequence_reset
+
+
+@pytest.fixture(autouse=True)
+def _dbdiff_marker(request):
+    marker = request.keywords.get('dbdiff', None)
+    if not marker:
+        return
+
+    # Enable transactional db
+    getfixturevalue(request, 'transactional_db')
+
+    for model in marker.kwargs['models']:
+        sequence_reset(model)
+
+
+def pytest_load_initial_conftests(early_config, parser, args):
+    """Register the dbdiff mark."""
+    early_config.addinivalue_line(
+        'markers',
+        'dbdiff(models, reset_sequences=True): Mark the test as using '
+        'the django test database.  The *transaction* argument marks will '
+        "allow you to use real transactions in the test like Django's "
+        'TransactionTestCase.')

dbdiff/sequence.py

@@ -0,0 +1,56 @@
+"""Smarter model pk sequence reset."""
+from django.db import connection, models
+
+
+def pk_sequence_get(model):
+    """Return a list of table, column tuples which should have sequences."""
+    for field in model._meta.get_fields():
+        if not getattr(field, 'primary_key', False):
+            continue
+        if not isinstance(field, models.AutoField):
+            continue
+        return field.db_column or field.column
+
+
+def sequence_reset(model):
+    """
+    Better sequence reset than TransactionTestCase.
+
+    The difference with using TransactionTestCase with reset_sequences=True is
+    that this will reset sequences for the given models to their higher value,
+    supporting pre-existing models which could have been created by a
+    migration.
+    """
+    pk_field = pk_sequence_get(model)
+    if not pk_field:
+        return
+
+    if connection.vendor == 'postgresql':
+        reset = """
+        SELECT
+            setval(
+                pg_get_serial_sequence('{table}', '{column}'),
+                coalesce(max({column}),0) + 1,
+                false
+            )
+        FROM {table}
+        """
+    elif connection.vendor == 'sqlite':
+        reset = """
+        UPDATE sqlite_sequence
+        SET seq=(SELECT max({column}) from {table})
+        WHERE name='{table}'
+        """
+    elif connection.vendor == 'mysql':
+        cursor = connection.cursor()
+        cursor.execute(
+            'SELECT MAX({column}) + 1 FROM {table}'.format(
+                column=pk_field, table=model._meta.db_table
+            )
+        )
+        result = cursor.fetchone()[0] or 0
+        reset = 'ALTER TABLE {table} AUTO_INCREMENT = %s' % result
+
+    connection.cursor().execute(
+        reset.format(column=pk_field, table=model._meta.db_table)
+    )

dbdiff/test.py

@@ -1,8 +1,8 @@
 """Convenience test mixin."""
 from django.core.management import call_command
-from django.db import connection, models
 
 from .fixture import Fixture
+from .sequence import sequence_reset
 
 
 class DbdiffTestMixin(object):
@@ -45,34 +45,7 @@ def test_db_import(self):
             call_command('loaddata', fixture)
 
         for model in self.dbdiff_models:
-            if connection.vendor == 'postgresql':
-                pk_field = None
-                for field in model._meta.get_fields():
-                    if getattr(field, 'primary_key', False):
-                        pk_field = field
-                        break
-
-                if not isinstance(pk_field, models.AutoField):
-                    continue
-
-                if not pk_field.db_column:
-                    continue
-
-                reset = """
-                SELECT
-                    setval(
-                        pg_get_serial_sequence('%(table)s', '%(column)s'),
-                        coalesce(max(%(column)s),0) + 1,
-                        false
-                    )
-                FROM %(table)s
-                """ % dict(
-                    table=model._meta.db_table,
-                    column=pk_field.db_column,
-                )
-            else:
-                raise NotImplemented()
-            connection.cursor().execute(reset)
+            sequence_reset(model)
 
         self.dbdiff_test()
 

dbdiff/tests/nonintpk/__init__.py

@@ -0,0 +1 @@
+"""Test that we don't crash with non sequence pks."""

dbdiff/tests/nonintpk/models.py

@@ -0,0 +1,8 @@
+import uuid
+
+from django.db import models
+
+
+class Nonintpk(models.Model):
+    # dbdiff should not try to reset this sequence
+    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)

dbdiff/tests/project/settings.py

@@ -40,6 +40,7 @@
     'dbdiff',
 
     'dbdiff.tests.decimal_test',
+    'dbdiff.tests.nonintpk',
 )
 
 MIDDLEWARE_CLASSES = (

dbdiff/tests/test_plugin.py

@@ -0,0 +1,21 @@
+from dbdiff.tests.decimal_test.models import TestModel as DecimalModel
+from dbdiff.tests.nonintpk.models import Nonintpk
+
+import pytest
+
+
+@pytest.mark.dbdiff(models=[DecimalModel])
+def test_insert_first():
+    assert DecimalModel.objects.count() == 0
+    assert DecimalModel.objects.create(test_field=1).pk == 1
+
+
+@pytest.mark.dbdiff(models=[DecimalModel])
+def test_still_first_pk():
+    assert DecimalModel.objects.count() == 0
+    assert DecimalModel.objects.create(test_field=1).pk == 1
+
+
+@pytest.mark.dbdiff(models=[DecimalModel, Nonintpk])
+def test_doesnt_reset_nonintpk_which_would_fail():
+    assert DecimalModel.objects.count() == 0

setup.py

@@ -23,6 +23,7 @@ def read(fname):
     license='MIT',
     keywords='django test database fixture diff',
     install_requires=['ijson', 'json_delta', 'six'],
+    entry_points={'pytest11': ['dbdiff = dbdiff.plugin']},
     classifiers=[
         'Development Status :: 4 - Beta',
         'Environment :: Web Environment',