Allow message assertions to check response context as well.

This makes it a little more flexible in edge cases where
the messages have already been unset in the cookie storage
(e.g. the request-response cycle is complete).

Change-Id: I9f3b1ec3f908d05c523ce013ab5fbd73837aff55
This commit is contained in:
Gabriel Hurley 2012-03-19 12:55:19 -07:00
parent 67f3d28349
commit 2f946fa07b

View File

@ -149,14 +149,14 @@ class TestCase(django_test.TestCase):
('Location', settings.TESTSERVER + expected_url)) ('Location', settings.TESTSERVER + expected_url))
self.assertEqual(response.status_code, 302) self.assertEqual(response.status_code, 302)
def assertNoMessages(self): def assertNoMessages(self, response=None):
""" """
Asserts that no messages have been attached by the ``contrib.messages`` Asserts that no messages have been attached by the ``contrib.messages``
framework. framework.
""" """
self.assertMessageCount(success=0, warn=0, info=0, error=0) self.assertMessageCount(response, success=0, warn=0, info=0, error=0)
def assertMessageCount(self, **kwargs): def assertMessageCount(self, response=None, **kwargs):
""" """
Asserts that the specified number of messages have been attached Asserts that the specified number of messages have been attached
for various message types. Usage would look like for various message types. Usage would look like
@ -166,10 +166,15 @@ class TestCase(django_test.TestCase):
temp_req.COOKIES = self.client.cookies temp_req.COOKIES = self.client.cookies
storage = default_storage(temp_req) storage = default_storage(temp_req)
messages = [] messages = []
# To gain early access to the messages we have to decode the
# cookie on the test client. if response is None:
if 'messages' in self.client.cookies: # To gain early access to the messages we have to decode the
messages = storage._decode(self.client.cookies['messages'].value) # cookie on the test client.
if 'messages' in self.client.cookies:
message_cookie = self.client.cookies['messages'].value
messages = storage._decode(message_cookie)
elif "messages" in response.context:
messages = response.context["messages"]
# If we don't have messages and we don't expect messages, we're done. # If we don't have messages and we don't expect messages, we're done.
if not any(kwargs.values()) and not messages: if not any(kwargs.values()) and not messages: