Replaced try..except blocks by context manager in custom lookups tests
This commit is contained in:
Andriy Sokolovskiy
Tim Graham
parent
0d71349773
commit
13dca01af0
|
|
@ -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'),
|
||||
["<Author: Foo>"], 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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue