bots: Make StateHandler store state on Zulip server.
This makes the StateHandler functional. To reduce the number of server roundtrips when fetching/updating the state, the entire state is fetched ocne at bot initialization and cached. All changes are stored in the cache and only saved externally after handle_message() has been executed. Fixes #141.
This commit is contained in:
parent
979bbb1c14
commit
2736223073
|
@ -14,7 +14,7 @@ from contextlib import contextmanager
|
||||||
|
|
||||||
if False:
|
if False:
|
||||||
from mypy_extensions import NoReturn
|
from mypy_extensions import NoReturn
|
||||||
from typing import Any, Optional, List, Dict, IO, Text
|
from typing import Any, Optional, List, Dict, IO, Text, Set
|
||||||
from types import ModuleType
|
from types import ModuleType
|
||||||
|
|
||||||
from zulip import Client
|
from zulip import Client
|
||||||
|
@ -52,16 +52,26 @@ class RateLimit(object):
|
||||||
logging.error(self.error_message)
|
logging.error(self.error_message)
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
|
class StateHandlerError(Exception):
|
||||||
|
pass
|
||||||
|
|
||||||
class StateHandler(object):
|
class StateHandler(object):
|
||||||
def __init__(self):
|
def __init__(self, client):
|
||||||
# type: () -> None
|
# type: (Client) -> None
|
||||||
self.state_ = {} # type: Dict[Text, Text]
|
self._client = client
|
||||||
self.marshal = lambda obj: json.dumps(obj)
|
self.marshal = lambda obj: json.dumps(obj)
|
||||||
self.demarshal = lambda obj: json.loads(obj)
|
self.demarshal = lambda obj: json.loads(obj)
|
||||||
|
response = self._client.get_state()
|
||||||
|
if response['result'] == 'success':
|
||||||
|
self.state_ = response['state']
|
||||||
|
self._modified_entries = set() # type: Set[Text]
|
||||||
|
else:
|
||||||
|
raise StateHandlerError("Error initializing state: {}".format(str(response)))
|
||||||
|
|
||||||
def put(self, key, value):
|
def put(self, key, value):
|
||||||
# type: (Text, Text) -> None
|
# type: (Text, Text) -> None
|
||||||
self.state_[key] = self.marshal(value)
|
self.state_[key] = self.marshal(value)
|
||||||
|
self._modified_entries.add(key)
|
||||||
|
|
||||||
def get(self, key):
|
def get(self, key):
|
||||||
# type: (Text) -> Text
|
# type: (Text) -> Text
|
||||||
|
@ -71,6 +81,16 @@ class StateHandler(object):
|
||||||
# type: (Text) -> bool
|
# type: (Text) -> bool
|
||||||
return key in self.state_
|
return key in self.state_
|
||||||
|
|
||||||
|
def _save(self):
|
||||||
|
# type: () -> None
|
||||||
|
state_update = {'state': {key: self.state_[key] for key in self._modified_entries}}
|
||||||
|
if state_update:
|
||||||
|
response = self._client.update_state(state_update)
|
||||||
|
if response['result'] == 'success':
|
||||||
|
self._modified_entries.clear()
|
||||||
|
else:
|
||||||
|
raise StateHandlerError("Error updating state: {}".format(str(response)))
|
||||||
|
|
||||||
class ExternalBotHandler(object):
|
class ExternalBotHandler(object):
|
||||||
def __init__(self, client, root_dir):
|
def __init__(self, client, root_dir):
|
||||||
# type: (Client, str) -> None
|
# type: (Client, str) -> None
|
||||||
|
@ -79,7 +99,7 @@ class ExternalBotHandler(object):
|
||||||
self._rate_limit = RateLimit(20, 5)
|
self._rate_limit = RateLimit(20, 5)
|
||||||
self._client = client
|
self._client = client
|
||||||
self._root_dir = root_dir
|
self._root_dir = root_dir
|
||||||
self.storage = StateHandler()
|
self.storage = StateHandler(client)
|
||||||
try:
|
try:
|
||||||
self.user_id = user_profile['user_id']
|
self.user_id = user_profile['user_id']
|
||||||
self.full_name = user_profile['full_name']
|
self.full_name = user_profile['full_name']
|
||||||
|
@ -218,6 +238,7 @@ def run_message_handler_for_bot(lib_module, quiet, config_file, bot_name):
|
||||||
message=message,
|
message=message,
|
||||||
bot_handler=restricted_client
|
bot_handler=restricted_client
|
||||||
)
|
)
|
||||||
|
restricted_client.storage._save()
|
||||||
|
|
||||||
signal.signal(signal.SIGINT, exit_gracefully)
|
signal.signal(signal.SIGINT, exit_gracefully)
|
||||||
|
|
||||||
|
|
|
@ -45,7 +45,9 @@ class BotTestCaseBase(TestCase):
|
||||||
self.patcher = patch('zulip_bots.lib.ExternalBotHandler', autospec=True)
|
self.patcher = patch('zulip_bots.lib.ExternalBotHandler', autospec=True)
|
||||||
self.MockClass = self.patcher.start()
|
self.MockClass = self.patcher.start()
|
||||||
self.mock_bot_handler = self.MockClass(None, None)
|
self.mock_bot_handler = self.MockClass(None, None)
|
||||||
self.mock_bot_handler.storage = StateHandler()
|
self.mock_client = MagicMock()
|
||||||
|
self.mock_client.get_state.return_value = {'result': 'success', 'state': {}}
|
||||||
|
self.mock_bot_handler.storage = StateHandler(self.mock_client)
|
||||||
self.mock_bot_handler.send_message.return_value = {'id': 42}
|
self.mock_bot_handler.send_message.return_value = {'id': 42}
|
||||||
self.mock_bot_handler.send_reply.return_value = {'id': 42}
|
self.mock_bot_handler.send_reply.return_value = {'id': 42}
|
||||||
self.message_handler = self.get_bot_message_handler()
|
self.message_handler = self.get_bot_message_handler()
|
||||||
|
|
|
@ -39,8 +39,9 @@ class BotServerTests(BotServerTestCase):
|
||||||
check_success=False)
|
check_success=False)
|
||||||
|
|
||||||
@mock.patch('logging.error')
|
@mock.patch('logging.error')
|
||||||
def test_wrong_bot_credentials(self, mock_LoggingError):
|
@mock.patch('zulip_bots.lib.StateHandler')
|
||||||
# type: (mock.Mock) -> None
|
def test_wrong_bot_credentials(self, mock_StateHandler, mock_LoggingError):
|
||||||
|
# type: (mock.Mock, mock.Mock) -> None
|
||||||
available_bots = ['nonexistent-bot']
|
available_bots = ['nonexistent-bot']
|
||||||
bots_config = {
|
bots_config = {
|
||||||
'nonexistent-bot': {
|
'nonexistent-bot': {
|
||||||
|
|
Loading…
Reference in a new issue