from django.contrib.messages.tests.base import BaseTest from django.contrib.messages.storage.session import SessionStorage def set_session_data(storage, messages): """ Sets the messages into the backend request's session and remove the backend's loaded data cache. """ storage.request.session[storage.session_key] = messages if hasattr(storage, '_loaded_data'): del storage._loaded_data def stored_session_messages_count(storage): data = storage.request.session.get(storage.session_key, []) return len(data) class SessionTest(BaseTest): storage_class = SessionStorage def get_request(self): self.session = {} request = super(SessionTest, self).get_request() request.session = self.session return request def stored_messages_count(self, storage, response): return stored_session_messages_count(storage) def test_get(self): storage = self.storage_class(self.get_request()) # Set initial data. example_messages = ['test', 'me'] set_session_data(storage, example_messages) # Test that the message actually contains what we expect. self.assertEqual(list(storage), example_messages)