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:
derAnfaenger 2017-11-01 13:26:20 +01:00 committed by showell
parent 979bbb1c14
commit 2736223073
3 changed files with 32 additions and 8 deletions

View file

@ -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)

View file

@ -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()

View file

@ -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': {