Added context manager dbdiff.exact

Author: jpic

README.rst

@@ -80,6 +80,12 @@ Example::
 If any difference is found between the database and the test fixture, then
 ``diff()`` will return the diff as outputed by GNU diff.
 
+A context manager that will raise an exception if a diff is found is also
+provided, it's able to delete models and reset the PK sequences for them::
+
+    with dbdiff.exact('your_app/fixture.json'):
+        # do stuff
+
 More public API tests can be found in dbidff/tests/test_dbdiff.py.
 
 Django model observer

dbdiff/dbdiff.py

@@ -1,14 +1,95 @@
 """Public API for dbdiff."""
 
 from django.apps import apps
+from django.core.management.color import no_style
+from django.db import connections
 
+from .exceptions import DiffFound
 from .fixture_diff import FixtureDiff
 
 __all__ = ('diff',)
 
 DEBUG = apps.get_app_config('dbdiff').debug
 
 
+class exact(object):  # noqa
+    """
+    Context manager interface for dbdiff.
+
+    To clean up the database from any model that's provided by the fixture on
+    context enter, and raise a :py:class:`~dbdiff.exceptions.DiffFound` on
+    context enter if any diff was found::
+
+        with dbdiff.exact('your_app/tests/your_test_expected.json'):
+            # do stuff
+
+    .. py:attribute:: diff
+
+        :py:class:`~dbdiff.fixture_diff.FixtureDiff` instance for fixture.
+
+    .. py:attribute:: database
+
+        Database to use, 'default' by default.
+
+    .. py:attribute:: reset_models
+
+        If False, do not delete models and reset PK sequences on context enter.
+    """
+
+    def __init__(self, fixture, reset_models=True, database=None):
+        """Instanciate with a fixture."""
+        self.diff = FixtureDiff(fixture)
+        self.database = database if database else 'default'
+        self.reset_models = reset_models
+
+    def __enter__(self):
+        if not self.reset_models:
+            return self
+
+        tables = set()
+
+        for model_name in self.diff.fixture_models:
+            model = apps.get_model(model_name)
+            tables.add(model._meta.db_table)
+            tables.update(f.m2m_db_table() for f in
+                          model._meta.local_many_to_many)
+
+        connection = connections[self.database]
+
+        statements = connection.ops.sql_flush(
+            no_style(),
+            list(tables),
+            connection.introspection.sequence_list(),
+            True
+        )
+        if connection.settings_dict['ENGINE'] == 'django.db.backends.sqlite3':
+            # Initially, we were just doing a model.objects.all().delete() in
+            # this method, and it worked only in SQLite:
+            # https://travis-ci.org/yourlabs/django-dbdiff/builds/100059628
+            # That's why this method now uses sqlflush method, but then PKs are
+            # not reset anymore in SQLite:
+            # http://stackoverflow.com/questions/24098733/why-doesnt-django-reset-sequences-in-sqlite3  # noqa
+            statements += [
+                "UPDATE SQLITE_SEQUENCE SET SEQ=0 WHERE NAME='%s';" % t
+                for t in tables
+            ]
+        cursor = connection.cursor()
+
+        for statement in statements:
+            cursor.execute(statement)
+
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        out = self.diff.get_diff()
+
+        if not out:
+            self.diff.clean()
+        else:
+            print(self.diff.cmd)
+            raise DiffFound(out)
+
+
 def diff(fixture):
     """
     Return the diff between the database and a fixture file as a string.
@@ -28,6 +109,8 @@ def diff(fixture):
     except:
         if not DEBUG:
             diff.clean()
+        else:
+            print(diff.cmd)
         raise
 
     if not DEBUG:

dbdiff/exceptions.py

@@ -11,3 +11,11 @@ class EmptyFixtures(DbDiffException):
     def __init__(self, path):
         """Exception for when path does not contain any fixture data."""
         super(EmptyFixtures, self).__init__('%s is empty' % path)
+
+
+class DiffFound(DbDiffException):
+    """Raised when a diff is found by the context manager."""
+
+    def __init__(self, out):
+        """Exception for when a diff command had output."""
+        super(DiffFound, self).__init__(out)

dbdiff/fixture_diff.py

@@ -70,13 +70,13 @@ def get_diff(self):
                          traceback=True, indent=self.fixture_indent,
                          stdout=f)
 
-        cmd = 'diff -u1 %s %s | sed "1,2 d"' % (
+        self.cmd = 'diff -u1 %s %s | sed "1,2 d"' % (
             self.fixture_path, self.dump_path)
 
         if apps.get_app_config('dbdiff').debug:  # pragma: no cover
-            print(cmd)
+            print(self.cmd)
 
-        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
+        proc = subprocess.Popen(self.cmd, stdout=subprocess.PIPE, shell=True)
         out, err = proc.communicate()
 
         return out

dbdiff/fixtures/dbdiff_test_group.json

@@ -0,0 +1,10 @@
+[
+{
+    "fields": {
+        "name": "initial_name",
+        "permissions": []
+    },
+    "model": "auth.group",
+    "pk": 1
+}
+]

dbdiff/tests/test_dbdiff.py

@@ -7,6 +7,7 @@
 from django import test
 from django.apps import apps
 from django.conf import settings
+from django.contrib.auth.models import Group
 
 import mock
 
@@ -15,6 +16,7 @@
 from . import base
 from .project.decimal_test.models import TestModel as DecimalTestModel
 from .. import dbdiff
+from ..exceptions import DiffFound
 
 
 @mock.patch('dbdiff.dbdiff.FixtureDiff')
@@ -77,6 +79,35 @@ def test_data_diff_has_changed(self):
 ''').lstrip()
 
 
+class ResetModelsTest(test.TestCase):
+    def test_pk_reset(self):
+        fixture = 'dbdiff/fixtures/dbdiff_test_group.json'
+        Group.objects.create(name='noise')
+
+        with dbdiff.exact(fixture):
+            Group.objects.create(name='initial_name')
+
+        with dbdiff.exact(fixture, reset_models=False):
+            pass
+
+        with self.assertRaises(DiffFound) as e:
+            with dbdiff.exact(fixture):
+                Group.objects.create(name='different_name')
+
+        expected = six.b('''
+@@ -3,3 +3,3 @@
+     "fields": {
+-        "name": "initial_name",
++        "name": "different_name",
+         "permissions": []
+''').lstrip()
+
+        if six.PY2:
+            assert e.exception.message == expected
+        else:
+            assert e.exception.args[0] == expected
+
+
 class DecimalDiffTest(test.TestCase):
     fixtures = ['decimal_test_fixture']
     expected = os.path.join(