|
9 | 9 | from rest_framework.exceptions import APIException |
10 | 10 | from rest_framework.response import Response |
11 | 11 | from rest_framework.test import APIRequestFactory |
12 | | -from rest_framework.views import APIView |
| 12 | +from rest_framework.views import APIView, set_rollback |
13 | 13 | from tests.models import BasicModel |
14 | 14 |
|
15 | 15 | factory = APIRequestFactory() |
@@ -183,3 +183,72 @@ def test_api_exception_rollback_transaction_non_atomic_view(self): |
183 | 183 | # without checking connection.in_atomic_block view raises 500 |
184 | 184 | # due attempt to rollback without transaction |
185 | 185 | assert response.status_code == status.HTTP_404_NOT_FOUND |
| 186 | + |
| 187 | + |
| 188 | +@unittest.skipUnless( |
| 189 | + connection.features.uses_savepoints, |
| 190 | + "'atomic' requires transactions and savepoints." |
| 191 | +) |
| 192 | +class SetRollbackTests(TestCase): |
| 193 | + def setUp(self): |
| 194 | + connections.databases['default']['ATOMIC_REQUESTS'] = True |
| 195 | + |
| 196 | + def tearDown(self): |
| 197 | + connections.databases['default']['ATOMIC_REQUESTS'] = False |
| 198 | + |
| 199 | + def test_marks_initialized_atomic_connection_for_rollback(self): |
| 200 | + with transaction.atomic(): |
| 201 | + set_rollback() |
| 202 | + assert transaction.get_rollback() |
| 203 | + |
| 204 | + |
| 205 | +class UninitializedSecondaryConnectionMixin: |
| 206 | + """ |
| 207 | + Remove the 'secondary' connection wrapper from the current thread for |
| 208 | + the duration of a test, restoring the original wrapper afterwards so |
| 209 | + Django's test-case connection patching still finds it at class cleanup. |
| 210 | + """ |
| 211 | + def setUp(self): |
| 212 | + super().setUp() |
| 213 | + self._saved_secondary = getattr(connections._connections, 'secondary', None) |
| 214 | + if self._saved_secondary is not None: |
| 215 | + delattr(connections._connections, 'secondary') |
| 216 | + |
| 217 | + def tearDown(self): |
| 218 | + stray = getattr(connections._connections, 'secondary', None) |
| 219 | + if stray is not None and stray is not self._saved_secondary: |
| 220 | + stray.close() |
| 221 | + if self._saved_secondary is not None: |
| 222 | + setattr(connections._connections, 'secondary', self._saved_secondary) |
| 223 | + super().tearDown() |
| 224 | + |
| 225 | + |
| 226 | +class SetRollbackUninitializedConnectionTests(UninitializedSecondaryConnectionMixin, TestCase): |
| 227 | + def setUp(self): |
| 228 | + super().setUp() |
| 229 | + connections.databases['secondary']['ATOMIC_REQUESTS'] = True |
| 230 | + |
| 231 | + def tearDown(self): |
| 232 | + connections.databases['secondary']['ATOMIC_REQUESTS'] = False |
| 233 | + super().tearDown() |
| 234 | + |
| 235 | + def test_does_not_initialize_unused_connections(self): |
| 236 | + set_rollback() |
| 237 | + assert not hasattr(connections._connections, 'secondary') |
| 238 | + |
| 239 | + |
| 240 | +class MultiDBUnusedConnectionAPIExceptionTests(UninitializedSecondaryConnectionMixin, TestCase): |
| 241 | + def setUp(self): |
| 242 | + super().setUp() |
| 243 | + self.view = APIExceptionView.as_view() |
| 244 | + connections.databases['secondary']['ATOMIC_REQUESTS'] = True |
| 245 | + |
| 246 | + def tearDown(self): |
| 247 | + connections.databases['secondary']['ATOMIC_REQUESTS'] = False |
| 248 | + super().tearDown() |
| 249 | + |
| 250 | + def test_api_exception_leaves_unused_connection_uninitialized(self): |
| 251 | + request = factory.post('/') |
| 252 | + response = self.view(request) |
| 253 | + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR |
| 254 | + assert not hasattr(connections._connections, 'secondary') |
0 commit comments