1 from django import http

     2 from django.contrib.auth.models import User

     3 from django.contrib.messages.storage.user_messages import UserMessagesStorage,\

     4     LegacyFallbackStorage

     5 from django.contrib.messages.tests.cookie import set_cookie_data

     6 from django.contrib.messages.tests.fallback import FallbackTest

     7 from django.test import TestCase

     8 

     9 

    10 class UserMessagesTest(TestCase):

    11 

    12     def setUp(self):

    13         self.user = User.objects.create(username='tester')

    14 

    15     def test_add(self):

    16         storage = UserMessagesStorage(http.HttpRequest())

    17         self.assertRaises(NotImplementedError, storage.add, 'Test message 1')

    18 

    19     def test_get_anonymous(self):

    20         # Ensure that the storage still works if no user is attached to the

    21         # request.

    22         storage = UserMessagesStorage(http.HttpRequest())

    23         self.assertEqual(len(storage), 0)

    24 

    25     def test_get(self):

    26         storage = UserMessagesStorage(http.HttpRequest())

    27         storage.request.user = self.user

    28         self.user.message_set.create(message='test message')

    29 

    30         self.assertEqual(len(storage), 1)

    31         self.assertEqual(list(storage)[0].message, 'test message')

    32 

    33 

    34 class LegacyFallbackTest(FallbackTest, TestCase):

    35     storage_class = LegacyFallbackStorage

    36 

    37     def setUp(self):

    38         super(LegacyFallbackTest, self).setUp()

    39         self.user = User.objects.create(username='tester')

    40 

    41     def get_request(self, *args, **kwargs):

    42         request = super(LegacyFallbackTest, self).get_request(*args, **kwargs)

    43         request.user = self.user

    44         return request

    45 

    46     def test_get_legacy_only(self):

    47         request = self.get_request()

    48         storage = self.storage_class(request)

    49         self.user.message_set.create(message='user message')

    50 

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

    52         self.assertEqual(len(storage), 1)

    53         self.assertEqual(list(storage)[0].message, 'user message')

    54 

    55     def test_get_legacy(self):

    56         request = self.get_request()

    57         storage = self.storage_class(request)

    58         cookie_storage = self.get_cookie_storage(storage)

    59         self.user.message_set.create(message='user message')

    60         set_cookie_data(cookie_storage, ['cookie'])

    61 

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

    63         self.assertEqual(len(storage), 2)

    64         self.assertEqual(list(storage)[0].message, 'user message')

    65         self.assertEqual(list(storage)[1], 'cookie')