Replaced try..except blocks by context manager in custom lookups tests

This commit is contained in:
Andriy Sokolovskiy 2015-07-03 19:24:00 +03:00 committed by Tim Graham
parent 0d71349773
commit 13dca01af0
1 changed files with 23 additions and 47 deletions

View File

@ -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)