Used CaptureQueriesContext in select_for_update tests.

This commit is contained in:
Simon Charette 2016-06-23 12:06:55 -04:00
parent d30febb4e5
commit cdf54db6c5
1 changed files with 8 additions and 11 deletions

View File

@ -10,13 +10,11 @@ from django.test import (
TransactionTestCase, override_settings, skipIfDBFeature, TransactionTestCase, override_settings, skipIfDBFeature,
skipUnlessDBFeature, skipUnlessDBFeature,
) )
from django.test.utils import CaptureQueriesContext
from .models import Person from .models import Person
# We need to set settings.DEBUG to True so we can capture the output SQL
# to examine.
@override_settings(DEBUG=True)
class SelectForUpdateTests(TransactionTestCase): class SelectForUpdateTests(TransactionTestCase):
available_apps = ['select_for_update'] available_apps = ['select_for_update']
@ -54,12 +52,11 @@ class SelectForUpdateTests(TransactionTestCase):
self.new_connection.rollback() self.new_connection.rollback()
self.new_connection.set_autocommit(True) self.new_connection.set_autocommit(True)
def has_for_update_sql(self, tested_connection, nowait=False): def has_for_update_sql(self, queries, nowait=False):
# Examine the SQL that was executed to determine whether it # Examine the SQL that was executed to determine whether it
# contains the 'SELECT..FOR UPDATE' stanza. # contains the 'SELECT..FOR UPDATE' stanza.
for_update_sql = tested_connection.ops.for_update_sql(nowait) for_update_sql = connection.ops.for_update_sql(nowait)
sql = tested_connection.queries[-1]['sql'] return any(for_update_sql in query['sql'] for query in queries)
return bool(sql.find(for_update_sql) > -1)
@skipUnlessDBFeature('has_select_for_update') @skipUnlessDBFeature('has_select_for_update')
def test_for_update_sql_generated(self): def test_for_update_sql_generated(self):
@ -67,9 +64,9 @@ class SelectForUpdateTests(TransactionTestCase):
Test that the backend's FOR UPDATE variant appears in Test that the backend's FOR UPDATE variant appears in
generated SQL when select_for_update is invoked. generated SQL when select_for_update is invoked.
""" """
with transaction.atomic(): with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
list(Person.objects.all().select_for_update()) list(Person.objects.all().select_for_update())
self.assertTrue(self.has_for_update_sql(connection)) self.assertTrue(self.has_for_update_sql(ctx.captured_queries))
@skipUnlessDBFeature('has_select_for_update_nowait') @skipUnlessDBFeature('has_select_for_update_nowait')
def test_for_update_sql_generated_nowait(self): def test_for_update_sql_generated_nowait(self):
@ -77,9 +74,9 @@ class SelectForUpdateTests(TransactionTestCase):
Test that the backend's FOR UPDATE NOWAIT variant appears in Test that the backend's FOR UPDATE NOWAIT variant appears in
generated SQL when select_for_update is invoked. generated SQL when select_for_update is invoked.
""" """
with transaction.atomic(): with transaction.atomic(), CaptureQueriesContext(connection) as ctx:
list(Person.objects.all().select_for_update(nowait=True)) list(Person.objects.all().select_for_update(nowait=True))
self.assertTrue(self.has_for_update_sql(connection, nowait=True)) self.assertTrue(self.has_for_update_sql(ctx.captured_queries, nowait=True))
@skipUnlessDBFeature('has_select_for_update_nowait') @skipUnlessDBFeature('has_select_for_update_nowait')
def test_nowait_raises_error_on_block(self): def test_nowait_raises_error_on_block(self):