Improve set_rollback() behaviour

rest_framework/views.py

@@ -3,7 +3,7 @@
 """
 from django.conf import settings
 from django.core.exceptions import PermissionDenied
-from django.db import connections, models
+from django.db import connections, models, transaction
 from django.http import Http404
 from django.http.response import HttpResponseBase
 from django.utils.cache import patch_vary_headers
@@ -64,9 +64,13 @@ def get_view_description(view, html=False):
 
 
 def set_rollback():
+    # Rollback all connections that have ATOMIC_REQUESTS set, if it looks like
+    # the @atomic block for the request was started.
+    # Note that this in_atomic_block check may be a false positive due to
+    # transactions started in other ways, e.g. when testing with TestCase.
     for db in connections.all(initialized_only=True):
         if db.settings_dict['ATOMIC_REQUESTS'] and db.in_atomic_block:
-            db.set_rollback(True)
+            transaction.set_rollback(True, using=db.alias)
 
 
 def exception_handler(exc, context):
@@ -230,7 +234,7 @@ def get_exception_handler_context(self):
             'view': self,
             'args': getattr(self, 'args', ()),
             'kwargs': getattr(self, 'kwargs', {}),
-            'request': getattr(self, 'request', None)
+            'request': getattr(self, 'request', None),
         }
 
     def get_view_name(self):

tests/test_atomic_requests.py

@@ -39,11 +39,12 @@ def dispatch(self, *args, **kwargs):
         return super().dispatch(*args, **kwargs)
 
     def get(self, request, *args, **kwargs):
-        BasicModel.objects.all()
+        list(BasicModel.objects.all())
         raise Http404
 
 
 urlpatterns = (
+    path('non-atomic-exception', NonAtomicAPIExceptionView.as_view()),
     path('', NonAtomicAPIExceptionView.as_view()),
 )
 
@@ -95,7 +96,8 @@ def test_generic_exception_delegate_transaction_management(self):
             # 2 - insert
             # 3 - release savepoint
             with transaction.atomic():
-                self.assertRaises(Exception, self.view, request)
+                with self.assertRaises(Exception):
+                    self.view(request)
                 assert not transaction.get_rollback()
         assert BasicModel.objects.count() == 1
 
@@ -174,15 +176,18 @@ class NonAtomicDBTransactionAPIExceptionTests(TransactionTestCase):
     def setUp(self):
         connections.databases['default']['ATOMIC_REQUESTS'] = True
 
-    def tearDown(self):
-        connections.databases['default']['ATOMIC_REQUESTS'] = False
+        @self.addCleanup
+        def restore_atomic_requests():
+            connections.databases['default']['ATOMIC_REQUESTS'] = False
 
     def test_api_exception_rollback_transaction_non_atomic_view(self):
-        response = self.client.get('/')
+        response = self.client.get('/non-atomic-exception')
 
-        # without checking connection.in_atomic_block view raises 500
-        # due attempt to rollback without transaction
+        # without check for db.in_atomic_block, would raise 500 due to attempt
+        # to rollback without transaction
         assert response.status_code == status.HTTP_404_NOT_FOUND
+        # Check we can still perform DB queries
+        list(BasicModel.objects.all())
 
 
 @unittest.skipUnless(