1. import random
    
  2. 
    
  3. from django.contrib.messages import constants
    
  4. from django.contrib.messages.storage.fallback import CookieStorage, FallbackStorage
    
  5. from django.test import SimpleTestCase
    
  6. from django.utils.crypto import get_random_string
    
  7. 
    
  8. from .base import BaseTests
    
  9. from .test_cookie import set_cookie_data, stored_cookie_messages_count
    
  10. from .test_session import set_session_data, stored_session_messages_count
    
  11. 
    
  12. 
    
  13. class FallbackTests(BaseTests, SimpleTestCase):
    
  14.     storage_class = FallbackStorage
    
  15. 
    
  16.     def get_request(self):
    
  17.         self.session = {}
    
  18.         request = super().get_request()
    
  19.         request.session = self.session
    
  20.         return request
    
  21. 
    
  22.     def get_cookie_storage(self, storage):
    
  23.         return storage.storages[-2]
    
  24. 
    
  25.     def get_session_storage(self, storage):
    
  26.         return storage.storages[-1]
    
  27. 
    
  28.     def stored_cookie_messages_count(self, storage, response):
    
  29.         return stored_cookie_messages_count(self.get_cookie_storage(storage), response)
    
  30. 
    
  31.     def stored_session_messages_count(self, storage, response):
    
  32.         return stored_session_messages_count(self.get_session_storage(storage))
    
  33. 
    
  34.     def stored_messages_count(self, storage, response):
    
  35.         """
    
  36.         Return the storage totals from both cookie and session backends.
    
  37.         """
    
  38.         return self.stored_cookie_messages_count(
    
  39.             storage, response
    
  40.         ) + self.stored_session_messages_count(storage, response)
    
  41. 
    
  42.     def test_get(self):
    
  43.         request = self.get_request()
    
  44.         storage = self.storage_class(request)
    
  45.         cookie_storage = self.get_cookie_storage(storage)
    
  46.         # Set initial cookie data.
    
  47.         example_messages = [str(i) for i in range(5)]
    
  48.         set_cookie_data(cookie_storage, example_messages)
    
  49.         # Overwrite the _get method of the fallback storage to prove it is not
    
  50.         # used (it would cause a TypeError: 'NoneType' object is not callable).
    
  51.         self.get_session_storage(storage)._get = None
    
  52.         self.assertEqual(list(storage), example_messages)
    
  53. 
    
  54.     def test_get_empty(self):
    
  55.         request = self.get_request()
    
  56.         storage = self.storage_class(request)
    
  57.         # Overwrite the _get method of the fallback storage to prove it is not
    
  58.         # used (it would cause a TypeError: 'NoneType' object is not callable).
    
  59.         self.get_session_storage(storage)._get = None
    
  60.         self.assertEqual(list(storage), [])
    
  61. 
    
  62.     def test_get_fallback(self):
    
  63.         request = self.get_request()
    
  64.         storage = self.storage_class(request)
    
  65.         cookie_storage = self.get_cookie_storage(storage)
    
  66.         session_storage = self.get_session_storage(storage)
    
  67. 
    
  68.         # Set initial cookie and session data.
    
  69.         example_messages = [str(i) for i in range(5)]
    
  70.         set_cookie_data(
    
  71.             cookie_storage, example_messages[:4] + [CookieStorage.not_finished]
    
  72.         )
    
  73.         set_session_data(session_storage, example_messages[4:])
    
  74.         self.assertEqual(list(storage), example_messages)
    
  75. 
    
  76.     def test_get_fallback_only(self):
    
  77.         request = self.get_request()
    
  78.         storage = self.storage_class(request)
    
  79.         cookie_storage = self.get_cookie_storage(storage)
    
  80.         session_storage = self.get_session_storage(storage)
    
  81.         # Set initial cookie and session data.
    
  82.         example_messages = [str(i) for i in range(5)]
    
  83.         set_cookie_data(cookie_storage, [CookieStorage.not_finished], encode_empty=True)
    
  84.         set_session_data(session_storage, example_messages)
    
  85.         self.assertEqual(list(storage), example_messages)
    
  86. 
    
  87.     def test_flush_used_backends(self):
    
  88.         request = self.get_request()
    
  89.         storage = self.storage_class(request)
    
  90.         cookie_storage = self.get_cookie_storage(storage)
    
  91.         session_storage = self.get_session_storage(storage)
    
  92.         # Set initial cookie and session data.
    
  93.         set_cookie_data(cookie_storage, ["cookie", CookieStorage.not_finished])
    
  94.         set_session_data(session_storage, ["session"])
    
  95.         # When updating, previously used but no longer needed backends are
    
  96.         # flushed.
    
  97.         response = self.get_response()
    
  98.         list(storage)
    
  99.         storage.update(response)
    
  100.         session_storing = self.stored_session_messages_count(storage, response)
    
  101.         self.assertEqual(session_storing, 0)
    
  102. 
    
  103.     def test_no_fallback(self):
    
  104.         """
    
  105.         (1) A short number of messages whose data size doesn't exceed what is
    
  106.         allowed in a cookie will all be stored in the CookieBackend.
    
  107. 
    
  108.         (2) If the CookieBackend can store all messages, the SessionBackend
    
  109.         won't be written to at all.
    
  110.         """
    
  111.         storage = self.get_storage()
    
  112.         response = self.get_response()
    
  113.         # Overwrite the _store method of the fallback storage to prove it isn't
    
  114.         # used (it would cause a TypeError: 'NoneType' object is not callable).
    
  115.         self.get_session_storage(storage)._store = None
    
  116.         for i in range(5):
    
  117.             storage.add(constants.INFO, str(i) * 100)
    
  118.         storage.update(response)
    
  119.         cookie_storing = self.stored_cookie_messages_count(storage, response)
    
  120.         self.assertEqual(cookie_storing, 5)
    
  121.         session_storing = self.stored_session_messages_count(storage, response)
    
  122.         self.assertEqual(session_storing, 0)
    
  123. 
    
  124.     def test_session_fallback(self):
    
  125.         """
    
  126.         If the data exceeds what is allowed in a cookie, messages which did
    
  127.         not fit are stored in the SessionBackend.
    
  128.         """
    
  129.         storage = self.get_storage()
    
  130.         response = self.get_response()
    
  131.         # see comment in CookieTests.test_cookie_max_length()
    
  132.         msg_size = int((CookieStorage.max_cookie_size - 54) / 4.5 - 37)
    
  133.         # Generate the same (tested) content every time that does not get run
    
  134.         # through zlib compression.
    
  135.         random.seed(42)
    
  136.         for i in range(5):
    
  137.             storage.add(constants.INFO, get_random_string(msg_size))
    
  138.         storage.update(response)
    
  139.         cookie_storing = self.stored_cookie_messages_count(storage, response)
    
  140.         self.assertEqual(cookie_storing, 4)
    
  141.         session_storing = self.stored_session_messages_count(storage, response)
    
  142.         self.assertEqual(session_storing, 1)
    
  143. 
    
  144.     def test_session_fallback_only(self):
    
  145.         """
    
  146.         Large messages, none of which fit in a cookie, are stored in the
    
  147.         SessionBackend (and nothing is stored in the CookieBackend).
    
  148.         """
    
  149.         storage = self.get_storage()
    
  150.         response = self.get_response()
    
  151.         # Generate the same (tested) content every time that does not get run
    
  152.         # through zlib compression.
    
  153.         random.seed(42)
    
  154.         storage.add(constants.INFO, get_random_string(5000))
    
  155.         storage.update(response)
    
  156.         cookie_storing = self.stored_cookie_messages_count(storage, response)
    
  157.         self.assertEqual(cookie_storing, 0)
    
  158.         session_storing = self.stored_session_messages_count(storage, response)
    
  159.         self.assertEqual(session_storing, 1)