Adjusted URLPatternsTestCase to use addClassCleanup() from Django 4.0.

carltongibson · carltongibson · commit 2d9eee5d022b · 2021-09-22T10:23:47.000+02:00
Refs django/django@faba5b7. addClassCleanup() is available from Python 3.8, which is the minimum supported Python from Django 4.0.
diff --git a/rest_framework/test.py b/rest_framework/test.py
@@ -3,6 +3,7 @@
 import io
 from importlib import import_module
 
+import django
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.core.handlers.wsgi import WSGIHandler
@@ -357,6 +358,13 @@ class APILiveServerTestCase(testcases.LiveServerTestCase):
     client_class = APIClient
 
 
+def cleanup_url_patterns(cls):
+    if hasattr(cls, '_module_urlpatterns'):
+        cls._module.urlpatterns = cls._module_urlpatterns
+    else:
+        del cls._module.urlpatterns
+
+
 class URLPatternsTestCase(testcases.SimpleTestCase):
     """
     Isolate URL patterns on a per-TestCase basis. For example,
@@ -385,14 +393,20 @@ def setUpClass(cls):
         cls._module.urlpatterns = cls.urlpatterns
 
         cls._override.enable()
+
+        if django.VERSION > (4, 0):
+            cls.addClassCleanup(cls._override.disable)
+            cls.addClassCleanup(cleanup_url_patterns, cls)
+
         super().setUpClass()
 
-    @classmethod
-    def tearDownClass(cls):
-        super().tearDownClass()
-        cls._override.disable()
+    if django.VERSION < (4, 0):
+        @classmethod
+        def tearDownClass(cls):
+            super().tearDownClass()
+            cls._override.disable()
 
-        if hasattr(cls, '_module_urlpatterns'):
-            cls._module.urlpatterns = cls._module_urlpatterns
-        else:
-            del cls._module.urlpatterns
+            if hasattr(cls, '_module_urlpatterns'):
+                cls._module.urlpatterns = cls._module_urlpatterns
+            else:
+                del cls._module.urlpatterns
diff --git a/tests/test_testing.py b/tests/test_testing.py
@@ -1,5 +1,6 @@
 from io import BytesIO
 
+import django
 from django.contrib.auth.models import User
 from django.shortcuts import redirect
 from django.test import TestCase, override_settings
@@ -282,6 +283,10 @@ def test_empty_request_content_type(self):
         assert request.META['CONTENT_TYPE'] == 'application/json'
 
 
+def check_urlpatterns(cls):
+    assert urlpatterns is not cls.urlpatterns
+
+
 class TestUrlPatternTestCase(URLPatternsTestCase):
     urlpatterns = [
         path('', view),
@@ -293,11 +298,18 @@ def setUpClass(cls):
         super().setUpClass()
         assert urlpatterns is cls.urlpatterns
 
-    @classmethod
-    def tearDownClass(cls):
-        assert urlpatterns is cls.urlpatterns
-        super().tearDownClass()
-        assert urlpatterns is not cls.urlpatterns
+        if django.VERSION > (4, 0):
+            cls.addClassCleanup(
+                check_urlpatterns,
+                cls
+            )
+
+    if django.VERSION < (4, 0):
+        @classmethod
+        def tearDownClass(cls):
+            assert urlpatterns is cls.urlpatterns
+            super().tearDownClass()
+            assert urlpatterns is not cls.urlpatterns
 
     def test_urlpatterns(self):
         assert self.client.get('/').status_code == 200
