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

1 files changed, 5 insertions, 24 deletions

diff --git a/tests/custom_lookups/tests.py b/tests/custom_lookups/tests.py
index 418525c3ed..4bf85339ed 100644
--- a/tests/custom_lookups/tests.py
+++ b/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):
         """