1 from django.contrib.messages.tests.base import BaseTest

     2 from django.contrib.messages.storage.session import SessionStorage

     3 

     4 

     5 def set_session_data(storage, messages):

     6     """

     7     Sets the messages into the backend request's session and remove the

     8     backend's loaded data cache.

     9     """

    10     storage.request.session[storage.session_key] = messages

    11     if hasattr(storage, '_loaded_data'):

    12         del storage._loaded_data

    13 

    14 

    15 def stored_session_messages_count(storage):

    16     data = storage.request.session.get(storage.session_key, [])

    17     return len(data)

    18 

    19 

    20 class SessionTest(BaseTest):

    21     storage_class = SessionStorage

    22 

    23     def get_request(self):

    24         self.session = {}

    25         request = super(SessionTest, self).get_request()

    26         request.session = self.session

    27         return request

    28 

    29     def stored_messages_count(self, storage, response):

    30         return stored_session_messages_count(storage)

    31 

    32     def test_get(self):

    33         storage = self.storage_class(self.get_request())

    34         # Set initial data.

    35         example_messages = ['test', 'me']

    36         set_session_data(storage, example_messages)

    37         # Test that the message actually contains what we expect.

    38         self.assertEqual(list(storage), example_messages)