Jelajahi Sumber

Fixed #29658 -- Registered model lookups in tests with a context manager.

Srinivas Reddy Thatiparthy (శ్రీనివాస్ రెడ్డి తాటిపర్తి) 6 tahun lalu
induk
melakukan
233c70f047

+ 15 - 0
django/test/utils.py

@@ -863,3 +863,18 @@ def tag(*tags):
             setattr(obj, 'tags', set(tags))
         return obj
     return decorator
+
+
+@contextmanager
+def register_lookup(field, *lookups, lookup_name=None):
+    """
+    Context manager to temporarily register lookups on a model field using
+    lookup_name (or the lookup's lookup_name if not provided).
+    """
+    try:
+        for lookup in lookups:
+            field.register_lookup(lookup, lookup_name)
+        yield
+    finally:
+        for lookup in lookups:
+            field._unregister_lookup(lookup, lookup_name)

+ 3 - 9
tests/admin_changelist/tests.py

@@ -16,7 +16,7 @@ from django.db.models.lookups import Contains, Exact
 from django.template import Context, Template, TemplateSyntaxError
 from django.test import TestCase, override_settings
 from django.test.client import RequestFactory
-from django.test.utils import CaptureQueriesContext
+from django.test.utils import CaptureQueriesContext, register_lookup
 from django.urls import reverse
 from django.utils import formats
 
@@ -480,8 +480,7 @@ class ChangeListTests(TestCase):
 
         m = ConcertAdmin(Concert, custom_site)
         m.search_fields = ['group__name__cc']
-        Field.register_lookup(Contains, 'cc')
-        try:
+        with register_lookup(Field, Contains, lookup_name='cc'):
             request = self.factory.get('/', data={SEARCH_VAR: 'Hype'})
             request.user = self.superuser
             cl = m.get_changelist_instance(request)
@@ -491,8 +490,6 @@ class ChangeListTests(TestCase):
             request.user = self.superuser
             cl = m.get_changelist_instance(request)
             self.assertCountEqual(cl.queryset, [])
-        finally:
-            Field._unregister_lookup(Contains, 'cc')
 
     def test_spanning_relations_with_custom_lookup_in_search_fields(self):
         hype = Group.objects.create(name='The Hype')
@@ -501,8 +498,7 @@ class ChangeListTests(TestCase):
         Membership.objects.create(music=vox, group=hype)
         # Register a custom lookup on IntegerField to ensure that field
         # traversing logic in ModelAdmin.get_search_results() works.
-        IntegerField.register_lookup(Exact, 'exactly')
-        try:
+        with register_lookup(IntegerField, Exact, lookup_name='exactly'):
             m = ConcertAdmin(Concert, custom_site)
             m.search_fields = ['group__members__age__exactly']
 
@@ -515,8 +511,6 @@ class ChangeListTests(TestCase):
             request.user = self.superuser
             cl = m.get_changelist_instance(request)
             self.assertCountEqual(cl.queryset, [])
-        finally:
-            IntegerField._unregister_lookup(Exact, 'exactly')
 
     def test_custom_lookup_with_pk_shortcut(self):
         self.assertEqual(CharPK._meta.pk.name, 'char_pk')  # Not equal to 'pk'.

+ 5 - 24
tests/custom_lookups/tests.py

@@ -1,4 +1,3 @@
-import contextlib
 import time
 import unittest
 from datetime import date, datetime
@@ -6,22 +5,12 @@ from datetime import date, datetime
 from django.core.exceptions import FieldError
 from django.db import connection, models
 from django.test import TestCase, override_settings
+from django.test.utils import register_lookup
 from django.utils import timezone
 
 from .models import Article, Author, MySQLUnixTimestamp
 
 
-@contextlib.contextmanager
-def register_lookup(field, *lookups):
-    try:
-        for lookup in lookups:
-            field.register_lookup(lookup)
-        yield
-    finally:
-        for lookup in lookups:
-            field._unregister_lookup(lookup)
-
-
 class Div3Lookup(models.Lookup):
     lookup_name = 'div3'
 
@@ -231,22 +220,14 @@ class LookupTests(TestCase):
     def test_custom_name_lookup(self):
         a1 = Author.objects.create(name='a1', birthdate=date(1981, 2, 16))
         Author.objects.create(name='a2', birthdate=date(2012, 2, 29))
-        custom_lookup_name = 'isactually'
-        custom_transform_name = 'justtheyear'
-        try:
-            models.DateField.register_lookup(YearTransform)
-            models.DateField.register_lookup(YearTransform, custom_transform_name)
-            YearTransform.register_lookup(Exactly)
-            YearTransform.register_lookup(Exactly, custom_lookup_name)
+        with register_lookup(models.DateField, YearTransform), \
+                register_lookup(models.DateField, YearTransform, lookup_name='justtheyear'), \
+                register_lookup(YearTransform, Exactly), \
+                register_lookup(YearTransform, Exactly, lookup_name='isactually'):
             qs1 = Author.objects.filter(birthdate__testyear__exactly=1981)
             qs2 = Author.objects.filter(birthdate__justtheyear__isactually=1981)
             self.assertSequenceEqual(qs1, [a1])
             self.assertSequenceEqual(qs2, [a1])
-        finally:
-            YearTransform._unregister_lookup(Exactly)
-            YearTransform._unregister_lookup(Exactly, custom_lookup_name)
-            models.DateField._unregister_lookup(YearTransform)
-            models.DateField._unregister_lookup(YearTransform, custom_transform_name)
 
     def test_custom_exact_lookup_none_rhs(self):
         """

+ 2 - 4
tests/db_functions/math/test_abs.py

@@ -3,6 +3,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Abs
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -40,11 +41,8 @@ class AbsTests(TestCase):
         self.assertEqual(obj.big, -obj.big_abs)
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Abs)
+        with register_lookup(DecimalField, Abs):
             DecimalModel.objects.create(n1=Decimal('-1.5'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('-0.5'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__abs__gt=1)
             self.assertQuerysetEqual(objs, [Decimal('-1.5')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Abs)

+ 2 - 4
tests/db_functions/math/test_acos.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import ACos
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class ACosTests(TestCase):
         self.assertAlmostEqual(obj.big_acos, math.acos(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(ACos)
+        with register_lookup(DecimalField, ACos):
             DecimalModel.objects.create(n1=Decimal('0.5'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('-0.9'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__acos__lt=2)
             self.assertQuerysetEqual(objs, [Decimal('0.5')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(ACos)

+ 2 - 4
tests/db_functions/math/test_asin.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import ASin
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class ASinTests(TestCase):
         self.assertAlmostEqual(obj.big_asin, math.asin(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(ASin)
+        with register_lookup(DecimalField, ASin):
             DecimalModel.objects.create(n1=Decimal('0.1'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('1.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__asin__gt=1)
             self.assertQuerysetEqual(objs, [Decimal('1.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(ASin)

+ 2 - 4
tests/db_functions/math/test_atan.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import ATan
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class ATanTests(TestCase):
         self.assertAlmostEqual(obj.big_atan, math.atan(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(ATan)
+        with register_lookup(DecimalField, ATan):
             DecimalModel.objects.create(n1=Decimal('3.12'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('-5'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__atan__gt=0)
             self.assertQuerysetEqual(objs, [Decimal('3.12')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(ATan)

+ 2 - 4
tests/db_functions/math/test_ceil.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Ceil
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class CeilTests(TestCase):
         self.assertEqual(obj.big_ceil, math.ceil(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Ceil)
+        with register_lookup(DecimalField, Ceil):
             DecimalModel.objects.create(n1=Decimal('3.12'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('1.25'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__ceil__gt=3)
             self.assertQuerysetEqual(objs, [Decimal('3.12')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Ceil)

+ 2 - 4
tests/db_functions/math/test_cos.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Cos
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class CosTests(TestCase):
         self.assertAlmostEqual(obj.big_cos, math.cos(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Cos)
+        with register_lookup(DecimalField, Cos):
             DecimalModel.objects.create(n1=Decimal('-8.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('3.14'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__cos__gt=-0.2)
             self.assertQuerysetEqual(objs, [Decimal('-8.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Cos)

+ 2 - 4
tests/db_functions/math/test_cot.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Cot
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class CotTests(TestCase):
         self.assertAlmostEqual(obj.big_cot, 1 / math.tan(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Cot)
+        with register_lookup(DecimalField, Cot):
             DecimalModel.objects.create(n1=Decimal('12.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('1.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__cot__gt=0)
             self.assertQuerysetEqual(objs, [Decimal('1.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Cot)

+ 2 - 4
tests/db_functions/math/test_degrees.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Degrees
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class DegreesTests(TestCase):
         self.assertAlmostEqual(obj.big_degrees, math.degrees(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Degrees)
+        with register_lookup(DecimalField, Degrees):
             DecimalModel.objects.create(n1=Decimal('5.4'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('-30'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__degrees__gt=0)
             self.assertQuerysetEqual(objs, [Decimal('5.4')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Degrees)

+ 2 - 4
tests/db_functions/math/test_exp.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Exp
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class ExpTests(TestCase):
         self.assertAlmostEqual(obj.big_exp, math.exp(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Exp)
+        with register_lookup(DecimalField, Exp):
             DecimalModel.objects.create(n1=Decimal('12.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('-1.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__exp__gt=10)
             self.assertQuerysetEqual(objs, [Decimal('12.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Exp)

+ 2 - 4
tests/db_functions/math/test_floor.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Floor
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class FloorTests(TestCase):
         self.assertEqual(obj.big_floor, math.floor(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Floor)
+        with register_lookup(DecimalField, Floor):
             DecimalModel.objects.create(n1=Decimal('5.4'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('3.4'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__floor__gt=4)
             self.assertQuerysetEqual(objs, [Decimal('5.4')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Floor)

+ 2 - 4
tests/db_functions/math/test_ln.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Ln
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class LnTests(TestCase):
         self.assertAlmostEqual(obj.big_ln, math.log(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Ln)
+        with register_lookup(DecimalField, Ln):
             DecimalModel.objects.create(n1=Decimal('12.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('1.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__ln__gt=0)
             self.assertQuerysetEqual(objs, [Decimal('12.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Ln)

+ 2 - 4
tests/db_functions/math/test_radians.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Radians
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class RadiansTests(TestCase):
         self.assertAlmostEqual(obj.big_radians, math.radians(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Radians)
+        with register_lookup(DecimalField, Radians):
             DecimalModel.objects.create(n1=Decimal('2.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('-1.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__radians__gt=0)
             self.assertQuerysetEqual(objs, [Decimal('2.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Radians)

+ 2 - 4
tests/db_functions/math/test_round.py

@@ -3,6 +3,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Round
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -40,11 +41,8 @@ class RoundTests(TestCase):
         self.assertEqual(obj.big_round, round(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Round)
+        with register_lookup(DecimalField, Round):
             DecimalModel.objects.create(n1=Decimal('2.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('-1.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__round__gt=0)
             self.assertQuerysetEqual(objs, [Decimal('2.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Round)

+ 2 - 4
tests/db_functions/math/test_sin.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Sin
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class SinTests(TestCase):
         self.assertAlmostEqual(obj.big_sin, math.sin(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Sin)
+        with register_lookup(DecimalField, Sin):
             DecimalModel.objects.create(n1=Decimal('5.4'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('0.1'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__sin__lt=0)
             self.assertQuerysetEqual(objs, [Decimal('5.4')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Sin)

+ 2 - 4
tests/db_functions/math/test_sqrt.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Sqrt
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class SqrtTests(TestCase):
         self.assertAlmostEqual(obj.big_sqrt, math.sqrt(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Sqrt)
+        with register_lookup(DecimalField, Sqrt):
             DecimalModel.objects.create(n1=Decimal('6.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('1.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__sqrt__gt=2)
             self.assertQuerysetEqual(objs, [Decimal('6.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Sqrt)

+ 2 - 4
tests/db_functions/math/test_tan.py

@@ -4,6 +4,7 @@ from decimal import Decimal
 from django.db.models import DecimalField
 from django.db.models.functions import Tan
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import DecimalModel, FloatModel, IntegerModel
 
@@ -41,11 +42,8 @@ class TanTests(TestCase):
         self.assertAlmostEqual(obj.big_tan, math.tan(obj.big))
 
     def test_transform(self):
-        try:
-            DecimalField.register_lookup(Tan)
+        with register_lookup(DecimalField, Tan):
             DecimalModel.objects.create(n1=Decimal('0.0'), n2=Decimal('0'))
             DecimalModel.objects.create(n1=Decimal('12.0'), n2=Decimal('0'))
             objs = DecimalModel.objects.filter(n1__tan__lt=0)
             self.assertQuerysetEqual(objs, [Decimal('12.0')], lambda a: a.n1)
-        finally:
-            DecimalField._unregister_lookup(Tan)

+ 7 - 14
tests/db_functions/tests.py

@@ -1,10 +1,15 @@
 from django.db.models import CharField, Value as V
 from django.db.models.functions import Coalesce, Length, Upper
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from .models import Author
 
 
+class UpperBilateral(Upper):
+    bilateral = True
+
+
 class FunctionTests(TestCase):
 
     def test_nested_function_ordering(self):
@@ -30,11 +35,7 @@ class FunctionTests(TestCase):
         )
 
     def test_func_transform_bilateral(self):
-        class UpperBilateral(Upper):
-            bilateral = True
-
-        try:
-            CharField.register_lookup(UpperBilateral)
+        with register_lookup(CharField, UpperBilateral):
             Author.objects.create(name='John Smith', alias='smithj')
             Author.objects.create(name='Rhonda')
             authors = Author.objects.filter(name__upper__exact='john smith')
@@ -44,15 +45,9 @@ class FunctionTests(TestCase):
                 ],
                 lambda a: a.name
             )
-        finally:
-            CharField._unregister_lookup(UpperBilateral)
 
     def test_func_transform_bilateral_multivalue(self):
-        class UpperBilateral(Upper):
-            bilateral = True
-
-        try:
-            CharField.register_lookup(UpperBilateral)
+        with register_lookup(CharField, UpperBilateral):
             Author.objects.create(name='John Smith', alias='smithj')
             Author.objects.create(name='Rhonda')
             authors = Author.objects.filter(name__upper__in=['john smith', 'rhonda'])
@@ -63,8 +58,6 @@ class FunctionTests(TestCase):
                 ],
                 lambda a: a.name
             )
-        finally:
-            CharField._unregister_lookup(UpperBilateral)
 
     def test_function_as_filter(self):
         Author.objects.create(name='John Smith', alias='SMITHJ')

+ 2 - 4
tests/db_functions/text/test_chr.py

@@ -1,6 +1,7 @@
 from django.db.models import IntegerField
 from django.db.models.functions import Chr, Left, Ord
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import Author
 
@@ -23,10 +24,7 @@ class ChrTests(TestCase):
         self.assertCountEqual(authors.exclude(first_initial=Chr(ord('É'))), [self.john, self.rhonda])
 
     def test_transform(self):
-        try:
-            IntegerField.register_lookup(Chr)
+        with register_lookup(IntegerField, Chr):
             authors = Author.objects.annotate(name_code_point=Ord('name'))
             self.assertCountEqual(authors.filter(name_code_point__chr=Chr(ord('J'))), [self.john])
             self.assertCountEqual(authors.exclude(name_code_point__chr=Chr(ord('J'))), [self.elena, self.rhonda])
-        finally:
-            IntegerField._unregister_lookup(Chr)

+ 2 - 4
tests/db_functions/text/test_length.py

@@ -1,6 +1,7 @@
 from django.db.models import CharField
 from django.db.models.functions import Length
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import Author
 
@@ -35,8 +36,7 @@ class LengthTests(TestCase):
         )
 
     def test_transform(self):
-        try:
-            CharField.register_lookup(Length)
+        with register_lookup(CharField, Length):
             Author.objects.create(name='John Smith', alias='smithj')
             Author.objects.create(name='Rhonda')
             authors = Author.objects.filter(name__length__gt=7)
@@ -44,5 +44,3 @@ class LengthTests(TestCase):
                 authors.order_by('name'), ['John Smith'],
                 lambda a: a.name
             )
-        finally:
-            CharField._unregister_lookup(Length)

+ 2 - 4
tests/db_functions/text/test_lower.py

@@ -1,6 +1,7 @@
 from django.db.models import CharField
 from django.db.models.functions import Lower
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import Author
 
@@ -29,8 +30,7 @@ class LowerTests(TestCase):
             Author.objects.update(name=Lower('name', 'name'))
 
     def test_transform(self):
-        try:
-            CharField.register_lookup(Lower)
+        with register_lookup(CharField, Lower):
             Author.objects.create(name='John Smith', alias='smithj')
             Author.objects.create(name='Rhonda')
             authors = Author.objects.filter(name__lower__exact='john smith')
@@ -38,5 +38,3 @@ class LowerTests(TestCase):
                 authors.order_by('name'), ['John Smith'],
                 lambda a: a.name
             )
-        finally:
-            CharField._unregister_lookup(Lower)

+ 2 - 4
tests/db_functions/text/test_ord.py

@@ -1,6 +1,7 @@
 from django.db.models import CharField, Value
 from django.db.models.functions import Left, Ord
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import Author
 
@@ -18,10 +19,7 @@ class OrdTests(TestCase):
         self.assertCountEqual(authors.exclude(name_part__gt=Ord(Value('John'))), [self.john])
 
     def test_transform(self):
-        try:
-            CharField.register_lookup(Ord)
+        with register_lookup(CharField, Ord):
             authors = Author.objects.annotate(first_initial=Left('name', 1))
             self.assertCountEqual(authors.filter(first_initial__ord=ord('J')), [self.john])
             self.assertCountEqual(authors.exclude(first_initial__ord=ord('J')), [self.elena, self.rhonda])
-        finally:
-            CharField._unregister_lookup(Ord)

+ 2 - 4
tests/db_functions/text/test_trim.py

@@ -1,6 +1,7 @@
 from django.db.models import CharField
 from django.db.models.functions import LTrim, RTrim, Trim
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import Author
 
@@ -32,9 +33,6 @@ class TrimTests(TestCase):
         )
         for transform, trimmed_name in tests:
             with self.subTest(transform=transform):
-                try:
-                    CharField.register_lookup(transform)
+                with register_lookup(CharField, transform):
                     authors = Author.objects.filter(**{'name__%s' % transform.lookup_name: trimmed_name})
                     self.assertQuerysetEqual(authors, [' John  '], lambda a: a.name)
-                finally:
-                    CharField._unregister_lookup(transform)

+ 2 - 4
tests/db_functions/text/test_upper.py

@@ -1,6 +1,7 @@
 from django.db.models import CharField
 from django.db.models.functions import Upper
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from ..models import Author
 
@@ -28,8 +29,7 @@ class UpperTests(TestCase):
         )
 
     def test_transform(self):
-        try:
-            CharField.register_lookup(Upper)
+        with register_lookup(CharField, Upper):
             Author.objects.create(name='John Smith', alias='smithj')
             Author.objects.create(name='Rhonda')
             authors = Author.objects.filter(name__upper__exact='JOHN SMITH')
@@ -39,5 +39,3 @@ class UpperTests(TestCase):
                 ],
                 lambda a: a.name
             )
-        finally:
-            CharField._unregister_lookup(Upper)

+ 2 - 4
tests/distinct_on_fields/tests.py

@@ -1,6 +1,7 @@
 from django.db.models import CharField, Max
 from django.db.models.functions import Lower
 from django.test import TestCase, skipUnlessDBFeature
+from django.test.utils import register_lookup
 
 from .models import Celebrity, Fan, Staff, StaffTag, Tag
 
@@ -100,14 +101,11 @@ class DistinctOnTests(TestCase):
         new_name = self.t1.name.upper()
         self.assertNotEqual(self.t1.name, new_name)
         Tag.objects.create(name=new_name)
-        CharField.register_lookup(Lower)
-        try:
+        with register_lookup(CharField, Lower):
             self.assertCountEqual(
                 Tag.objects.order_by().distinct('name__lower'),
                 [self.t1, self.t2, self.t3, self.t4, self.t5],
             )
-        finally:
-            CharField._unregister_lookup(Lower)
 
     def test_distinct_not_implemented_checks(self):
         # distinct + annotate not allowed

+ 2 - 4
tests/queries/test_query.py

@@ -9,6 +9,7 @@ from django.db.models.lookups import Exact, GreaterThan, IsNull, LessThan
 from django.db.models.sql.query import Query
 from django.db.models.sql.where import OR
 from django.test import TestCase
+from django.test.utils import register_lookup
 
 from .models import Author, Item, ObjectC, Ranking
 
@@ -49,11 +50,8 @@ class TestQuery(TestCase):
 
     def test_transform(self):
         query = Query(Author)
-        CharField.register_lookup(Lower, 'lower')
-        try:
+        with register_lookup(CharField, Lower):
             where = query.build_where(~Q(name__lower='foo'))
-        finally:
-            CharField._unregister_lookup(Lower, 'lower')
         lookup = where.children[0]
         self.assertIsInstance(lookup, Exact)
         self.assertIsInstance(lookup.lhs, Lower)