1 files changed, 10 insertions, 23 deletions
diff --git a/tests/backends/tests.py b/tests/backends/tests.py
index 488f8d518b..b196133a68 100644
--- a/tests/backends/tests.py
+++ b/tests/backends/tests.py
@@ -640,8 +640,7 @@ class FkConstraintsTests(TransactionTestCase):
         """
         When constraint checks are disabled, should be able to write bad data without IntegrityErrors.
         """
-        transaction.set_autocommit(False)
-        try:
+        with transaction.atomic():
             # Create an Article.
             models.Article.objects.create(headline="Test article", pub_date=datetime.datetime(2010, 9, 4), reporter=self.r)
             # Retrive it from the DB
@@ -653,17 +652,13 @@ class FkConstraintsTests(TransactionTestCase):
                 connection.enable_constraint_checking()
             except IntegrityError:
                 self.fail("IntegrityError should not have occurred.")
-            finally:
-                transaction.rollback()
-        finally:
-            transaction.set_autocommit(True)
+            transaction.set_rollback(True)
 
     def test_disable_constraint_checks_context_manager(self):
         """
         When constraint checks are disabled (using context manager), should be able to write bad data without IntegrityErrors.
         """
-        transaction.set_autocommit(False)
-        try:
+        with transaction.atomic():
             # Create an Article.
             models.Article.objects.create(headline="Test article", pub_date=datetime.datetime(2010, 9, 4), reporter=self.r)
             # Retrive it from the DB
@@ -674,31 +669,23 @@ class FkConstraintsTests(TransactionTestCase):
                     a.save()
             except IntegrityError:
                 self.fail("IntegrityError should not have occurred.")
-            finally:
-                transaction.rollback()
-        finally:
-            transaction.set_autocommit(True)
+            transaction.set_rollback(True)
 
     def test_check_constraints(self):
         """
         Constraint checks should raise an IntegrityError when bad data is in the DB.
         """
-        try:
-            transaction.set_autocommit(False)
+        with transaction.atomic():
             # Create an Article.
             models.Article.objects.create(headline="Test article", pub_date=datetime.datetime(2010, 9, 4), reporter=self.r)
             # Retrive it from the DB
             a = models.Article.objects.get(headline="Test article")
             a.reporter_id = 30
-            try:
-                with connection.constraint_checks_disabled():
-                    a.save()
-                    with self.assertRaises(IntegrityError):
-                        connection.check_constraints()
-            finally:
-                transaction.rollback()
-        finally:
-            transaction.set_autocommit(True)
+            with connection.constraint_checks_disabled():
+                a.save()
+                with self.assertRaises(IntegrityError):
+                    connection.check_constraints()
+            transaction.set_rollback(True)
 
 
 class ThreadTests(TestCase):