Fixed #17604 - Added context-manager capability to assertTemplateUsed and assertTemplateNotUsed. Thanks Greg Müllegger.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@17412 bcc190cf-cafb-0310-a4f2-bffc1f526a37

1 files changed, 75 insertions, 5 deletions

diff --git a/django/test/testcases.py b/django/test/testcases.py
index 53ea02a4f0..2e011e4cea 100644
--- a/django/test/testcases.py
+++ b/django/test/testcases.py
@@ -3,6 +3,7 @@ from __future__ import with_statement
 import os
 import re
 import sys
+from copy import copy
 from functools import wraps
 from urlparse import urlsplit, urlunsplit
 from xml.dom.minidom import parseString, Node
@@ -28,8 +29,10 @@ from django.forms.fields import CharField
 from django.http import QueryDict
 from django.test import _doctest as doctest
 from django.test.client import Client
+from django.test.signals import template_rendered
 from django.test.utils import (get_warnings_state, restore_warnings_state,
     override_settings)
+from django.test.utils import ContextList
 from django.utils import simplejson, unittest as ut2
 from django.utils.encoding import smart_str, force_unicode
 from django.views.static import serve
@@ -260,8 +263,53 @@ class _AssertNumQueriesContext(object):
         )
 
 
-class SimpleTestCase(ut2.TestCase):
+class _AssertTemplateUsedContext(object):
+    def __init__(self, test_case, template_name):
+        self.test_case = test_case
+        self.template_name = template_name
+        self.rendered_templates = []
+        self.rendered_template_names = []
+        self.context = ContextList()
+
+    def on_template_render(self, sender, signal, template, context, **kwargs):
+        self.rendered_templates.append(template)
+        self.rendered_template_names.append(template.name)
+        self.context.append(copy(context))
+
+    def test(self):
+        return self.template_name in self.rendered_template_names
+
+    def message(self):
+        return u'%s was not rendered.' % self.template_name
+
+    def __enter__(self):
+        template_rendered.connect(self.on_template_render)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        template_rendered.disconnect(self.on_template_render)
+        if exc_type is not None:
+            return
 
+        if not self.test():
+            message = self.message()
+            if len(self.rendered_templates) == 0:
+                message += u' No template was rendered.'
+            else:
+                message += u' Following templates were rendered: %s' % (
+                    ', '.join(self.rendered_template_names))
+            self.test_case.fail(message)
+
+
+class _AssertTemplateNotUsedContext(_AssertTemplateUsedContext):
+    def test(self):
+        return self.template_name not in self.rendered_template_names
+
+    def message(self):
+        return u'%s was rendered.' % self.template_name
+
+
+class SimpleTestCase(ut2.TestCase):
     def save_warnings_state(self):
         """
         Saves the state of the warnings module
@@ -612,14 +660,25 @@ class TransactionTestCase(SimpleTestCase):
             self.fail(msg_prefix + "The form '%s' was not used to render the"
                       " response" % form)
 
-    def assertTemplateUsed(self, response, template_name, msg_prefix=''):
+    def assertTemplateUsed(self, response=None, template_name=None, msg_prefix=''):
         """
         Asserts that the template with the provided name was used in rendering
-        the response.
+        the response. Also useable as context manager.
         """
+        if response is None and template_name is None:
+            raise TypeError(u'response and/or template_name argument must be provided')
+
         if msg_prefix:
             msg_prefix += ": "
 
+        # use assertTemplateUsed as context manager
+        if not hasattr(response, 'templates') or (response is None and template_name):
+            if response:
+                template_name = response
+                response = None
+            context = _AssertTemplateUsedContext(self, template_name)
+            return context
+
         template_names = [t.name for t in response.templates]
         if not template_names:
             self.fail(msg_prefix + "No templates used to render the response")
@@ -628,14 +687,25 @@ class TransactionTestCase(SimpleTestCase):
             " the response. Actual template(s) used: %s" %
                 (template_name, u', '.join(template_names)))
 
-    def assertTemplateNotUsed(self, response, template_name, msg_prefix=''):
+    def assertTemplateNotUsed(self, response=None, template_name=None, msg_prefix=''):
         """
         Asserts that the template with the provided name was NOT used in
-        rendering the response.
+        rendering the response. Also useable as context manager.
         """
+        if response is None and template_name is None:
+            raise TypeError(u'response and/or template_name argument must be provided')
+
         if msg_prefix:
             msg_prefix += ": "
 
+        # use assertTemplateUsed as context manager
+        if not hasattr(response, 'templates') or (response is None and template_name):
+            if response:
+                template_name = response
+                response = None
+            context = _AssertTemplateNotUsedContext(self, template_name)
+            return context
+
         template_names = [t.name for t in response.templates]
         self.assertFalse(template_name in template_names,
             msg_prefix + "Template '%s' was used unexpectedly in rendering"