From 13dca01af02e91f25ae1d96b09b857993714c96d Mon Sep 17 00:00:00 2001
From: Andriy Sokolovskiy <me@asokolovskiy.com>
Date: Fri, 3 Jul 2015 19:24:00 +0300
Subject: [PATCH] Replaced try..except blocks by context manager in custom
lookups tests
---
tests/custom_lookups/tests.py | 70 ++++++++++++-----------------------
1 file changed, 23 insertions(+), 47 deletions(-)
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'),
["<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)