diff --git a/zerver/lib/bot_storage.py b/zerver/lib/bot_storage.py index 47c065b379..d9a0adf5ad 100644 --- a/zerver/lib/bot_storage.py +++ b/zerver/lib/bot_storage.py @@ -1,3 +1,4 @@ +from django.conf import settings from django.db.models import Sum from django.db.models.query import F from django.db.models.functions import Length @@ -26,9 +27,7 @@ def get_bot_state_size(bot_profile, key=None): def set_bot_state(bot_profile, key, value): # type: (UserProfile, Text, Text) -> None - from zerver.lib.bot_lib import StateHandler - state_size_limit = StateHandler.state_size_limit - + state_size_limit = settings.USER_STATE_SIZE_LIMIT old_entry_size = get_bot_state_size(bot_profile, key) new_entry_size = len(key) + len(value) old_state_size = get_bot_state_size(bot_profile) diff --git a/zerver/tests/test_service_bot_system.py b/zerver/tests/test_service_bot_system.py index 1bae8ac034..f61cff1b0e 100644 --- a/zerver/tests/test_service_bot_system.py +++ b/zerver/tests/test_service_bot_system.py @@ -4,6 +4,9 @@ import json import mock from typing import Any, Union, Mapping, Callable +from django.conf import settings +from django.test import override_settings + from zerver.lib.actions import ( do_create_user, get_service_bot_events, @@ -146,23 +149,26 @@ class TestServiceBotStateHandler(ZulipTestCase): ", but it should be str."): storage.put(serializable_obj, 'some value') # type: ignore # We intend to test an invalid type. + # Reduce maximal state size for faster test string construction. + @override_settings(USER_STATE_SIZE_LIMIT=100) def test_storage_limit(self): # type: () -> None - # Reduce maximal state size for faster test string construction. - StateHandler.state_size_limit = 100 storage = StateHandler(self.bot_profile) - # Disable marshaling for storing a string whose size is equivalent to the size of the stored object. + + # Disable marshaling for storing a string whose size is + # equivalent to the size of the stored object. storage.marshal = lambda obj: obj storage.demarshal = lambda obj: obj + key = 'capacity-filling entry' - storage.put(key, 'x' * (StateHandler.state_size_limit - len(key))) + storage.put(key, 'x' * (settings.USER_STATE_SIZE_LIMIT - len(key))) with self.assertRaisesMessage(StateError, "Cannot set state. Request would require 132 bytes storage. " "The current storage limit is 100."): storage.put('too much data', 'a few bits too long') second_storage = StateHandler(self.second_bot_profile) - second_storage.put('another big entry', 'x' * (StateHandler.state_size_limit - 40)) + second_storage.put('another big entry', 'x' * (settings.USER_STATE_SIZE_LIMIT - 40)) second_storage.put('normal entry', 'abcd') def test_entry_removal(self): diff --git a/zproject/settings.py b/zproject/settings.py index c72c8a68cd..d2997f8394 100644 --- a/zproject/settings.py +++ b/zproject/settings.py @@ -159,6 +159,10 @@ DEFAULT_SETTINGS = { 'ENABLE_FEEDBACK': PRODUCTION, 'FEEDBACK_EMAIL': None, + # Max state storage per user + # TODO: Add this to zproject/prod_settings_template.py once stateful bots are fully functional. + 'USER_STATE_SIZE_LIMIT': 10000000, + # External service configuration 'CAMO_URI': '', 'MEMCACHED_LOCATION': '127.0.0.1:11211',