diff --git a/tests/custom_lookups/tests.py b/tests/custom_lookups/tests.py index dd8d431d3d..d8852e1e3c 100644 --- a/tests/custom_lookups/tests.py +++ b/tests/custom_lookups/tests.py @@ -1,5 +1,6 @@ from __future__ import unicode_literals +import contextlib import time import unittest from datetime import date, datetime @@ -12,6 +13,17 @@ from django.utils import timezone from .models import 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' @@ -193,8 +205,7 @@ class LookupTests(TestCase): a2 = Author.objects.create(name='a2', age=2) a3 = Author.objects.create(name='a3', age=3) a4 = Author.objects.create(name='a4', age=4) - models.IntegerField.register_lookup(Div3Lookup) - try: + with register_lookup(models.IntegerField, Div3Lookup): self.assertQuerysetEqual( Author.objects.filter(age__div3=0), [a3], lambda x: x @@ -211,8 +222,6 @@ class LookupTests(TestCase): Author.objects.filter(age__div3=3), [], lambda x: x ) - finally: - models.IntegerField._unregister_lookup(Div3Lookup) @unittest.skipUnless(connection.vendor == 'postgresql', "PostgreSQL specific SQL used") def test_birthdate_month(self): @@ -220,8 +229,7 @@ class LookupTests(TestCase): a2 = Author.objects.create(name='a2', birthdate=date(2012, 2, 29)) a3 = Author.objects.create(name='a3', birthdate=date(2012, 1, 31)) a4 = Author.objects.create(name='a4', birthdate=date(2012, 3, 1)) - models.DateField.register_lookup(InMonth) - try: + with register_lookup(models.DateField, InMonth): self.assertQuerysetEqual( Author.objects.filter(birthdate__inmonth=date(2012, 1, 15)), [a3], lambda x: x @@ -242,12 +250,9 @@ class LookupTests(TestCase): Author.objects.filter(birthdate__inmonth=date(2012, 4, 1)), [], lambda x: x ) - finally: - models.DateField._unregister_lookup(InMonth) def test_div3_extract(self): - models.IntegerField.register_lookup(Div3Transform) - try: + with register_lookup(models.IntegerField, Div3Transform): a1 = Author.objects.create(name='a1', age=1) a2 = Author.objects.create(name='a2', age=2) a3 = Author.objects.create(name='a3', age=3) @@ -271,15 +276,12 @@ class LookupTests(TestCase): self.assertQuerysetEqual( baseqs.filter(age__div3__range=(1, 2)), [a1, a2, a4], lambda x: x) - finally: - models.IntegerField._unregister_lookup(Div3Transform) class BilateralTransformTests(TestCase): def test_bilateral_upper(self): - models.CharField.register_lookup(UpperBilateralTransform) - try: + with register_lookup(models.CharField, UpperBilateralTransform): Author.objects.bulk_create([ Author(name='Doe'), Author(name='doe'), @@ -291,20 +293,14 @@ class BilateralTransformTests(TestCase): self.assertQuerysetEqual( Author.objects.filter(name__upper__contains='f'), [""], ordered=False) - finally: - models.CharField._unregister_lookup(UpperBilateralTransform) def test_bilateral_inner_qs(self): - models.CharField.register_lookup(UpperBilateralTransform) - try: + with register_lookup(models.CharField, UpperBilateralTransform): with self.assertRaises(NotImplementedError): Author.objects.filter(name__upper__in=Author.objects.values_list('name')) - finally: - models.CharField._unregister_lookup(UpperBilateralTransform) def test_div3_bilateral_extract(self): - models.IntegerField.register_lookup(Div3BilateralTransform) - try: + with register_lookup(models.IntegerField, Div3BilateralTransform): a1 = Author.objects.create(name='a1', age=1) a2 = Author.objects.create(name='a2', age=2) a3 = Author.objects.create(name='a3', age=3) @@ -328,13 +324,9 @@ class BilateralTransformTests(TestCase): self.assertQuerysetEqual( baseqs.filter(age__div3__range=(1, 2)), [a1, a2, a4], lambda x: x) - finally: - models.IntegerField._unregister_lookup(Div3BilateralTransform) def test_bilateral_order(self): - models.IntegerField.register_lookup(Mult3BilateralTransform) - models.IntegerField.register_lookup(Div3BilateralTransform) - try: + with register_lookup(models.IntegerField, Mult3BilateralTransform, Div3BilateralTransform): a1 = Author.objects.create(name='a1', age=1) a2 = Author.objects.create(name='a2', age=2) a3 = Author.objects.create(name='a3', age=3) @@ -348,13 +340,9 @@ class BilateralTransformTests(TestCase): self.assertQuerysetEqual( baseqs.filter(age__div3__mult3=42), [a3], lambda x: x) - finally: - models.IntegerField._unregister_lookup(Mult3BilateralTransform) - models.IntegerField._unregister_lookup(Div3BilateralTransform) def test_bilateral_fexpr(self): - models.IntegerField.register_lookup(Mult3BilateralTransform) - try: + with register_lookup(models.IntegerField, Mult3BilateralTransform): a1 = Author.objects.create(name='a1', age=1, average_rating=3.2) a2 = Author.objects.create(name='a2', age=2, average_rating=0.5) a3 = Author.objects.create(name='a3', age=3, average_rating=1.5) @@ -367,23 +355,18 @@ class BilateralTransformTests(TestCase): # Same as age >= average_rating baseqs.filter(age__mult3__gte=models.F('average_rating')), [a2, a3], lambda x: x) - finally: - models.IntegerField._unregister_lookup(Mult3BilateralTransform) @override_settings(USE_TZ=True) class DateTimeLookupTests(TestCase): @unittest.skipUnless(connection.vendor == 'mysql', "MySQL specific SQL used") def test_datetime_output_field(self): - models.PositiveIntegerField.register_lookup(DateTimeTransform) - try: + with register_lookup(models.PositiveIntegerField, DateTimeTransform): ut = MySQLUnixTimestamp.objects.create(timestamp=time.time()) y2k = timezone.make_aware(datetime(2000, 1, 1)) self.assertQuerysetEqual( MySQLUnixTimestamp.objects.filter(timestamp__as_datetime__gt=y2k), [ut], lambda x: x) - finally: - models.PositiveIntegerField._unregister_lookup(DateTimeTransform) class YearLteTests(TestCase): @@ -514,8 +497,7 @@ class TrackCallsYearTransform(YearTransform): class LookupTransformCallOrderTests(TestCase): def test_call_order(self): - models.DateField.register_lookup(TrackCallsYearTransform) - try: + with register_lookup(models.DateField, TrackCallsYearTransform): # junk lookup - tries lookup, then transform, then fails with self.assertRaises(FieldError): Author.objects.filter(birthdate__testyear__junk=2012) @@ -538,9 +520,6 @@ class LookupTransformCallOrderTests(TestCase): self.assertEqual(TrackCallsYearTransform.call_order, ['lookup']) - finally: - models.DateField._unregister_lookup(TrackCallsYearTransform) - class CustomisedMethodsTests(TestCase): @@ -563,8 +542,7 @@ class CustomisedMethodsTests(TestCase): class SubqueryTransformTests(TestCase): def test_subquery_usage(self): - models.IntegerField.register_lookup(Div3Transform) - try: + with register_lookup(models.IntegerField, Div3Transform): Author.objects.create(name='a1', age=1) a2 = Author.objects.create(name='a2', age=2) Author.objects.create(name='a3', age=3) @@ -572,5 +550,3 @@ class SubqueryTransformTests(TestCase): self.assertQuerysetEqual( Author.objects.order_by('name').filter(id__in=Author.objects.filter(age__div3=2)), [a2], lambda x: x) - finally: - models.IntegerField._unregister_lookup(Div3Transform)