From e526d0c14462f980a70e162167baf50b7a885415 Mon Sep 17 00:00:00 2001 From: derAnfaenger Date: Mon, 20 Nov 2017 14:40:51 +0100 Subject: [PATCH] embedded bots: Add views to access state. --- zerver/lib/bot_storage.py | 4 ++ zerver/tests/test_service_bot_system.py | 83 +++++++++++++++++++++++++ zerver/views/state.py | 45 ++++++++++++++ zproject/urls.py | 6 ++ 4 files changed, 138 insertions(+) create mode 100644 zerver/views/state.py diff --git a/zerver/lib/bot_storage.py b/zerver/lib/bot_storage.py index 7122c77efa..276b62aad1 100644 --- a/zerver/lib/bot_storage.py +++ b/zerver/lib/bot_storage.py @@ -57,3 +57,7 @@ def remove_bot_state(bot_profile, keys): def is_key_in_bot_state(bot_profile, key): # type: (UserProfile, Text) -> bool return BotUserStateData.objects.filter(bot_profile=bot_profile, key=key).exists() + +def get_keys_in_bot_state(bot_profile): + # type: (UserProfile) -> List[Text] + return list(BotUserStateData.objects.filter(bot_profile=bot_profile).values_list('key', flat=True)) diff --git a/zerver/tests/test_service_bot_system.py b/zerver/tests/test_service_bot_system.py index 26ac7fed65..11ad4ec535 100644 --- a/zerver/tests/test_service_bot_system.py +++ b/zerver/tests/test_service_bot_system.py @@ -22,6 +22,8 @@ from zerver.models import ( Recipient, ) +import ujson + BOT_TYPE_TO_QUEUE_NAME = { UserProfile.OUTGOING_WEBHOOK_BOT: 'outgoing_webhooks', UserProfile.EMBEDDED_BOT: 'embedded_bots', @@ -174,6 +176,87 @@ class TestServiceBotStateHandler(ZulipTestCase): self.assertTrue(storage.contains('another key')) self.assertRaises(StateError, lambda: storage.remove('some key')) + def test_internal_endpoint(self): + # type: () -> None + self.login(self.user_profile.email) + # Store some data. + initial_dict = {'key 1': 'value 1', 'key 2': 'value 2', 'key 3': 'value 3'} + params = { + 'state': ujson.dumps(initial_dict) + } + result = self.client_put('/json/user_state', params) + self.assert_json_success(result) + # Assert the stored data for some keys. + params = { + 'keys': ujson.dumps(['key 1', 'key 3']) + } + result = self.client_get('/json/user_state', params) + self.assert_json_success(result) + self.assertEqual(result.json()['state'], {'key 3': 'value 3', 'key 1': 'value 1'}) + # Assert the stored data for all keys. + result = self.client_get('/json/user_state') + self.assert_json_success(result) + self.assertEqual(result.json()['state'], initial_dict) + # Store some more data; update an entry and store a new entry + dict_update = {'key 1': 'new value', 'key 4': 'value 4'} + params = { + 'state': ujson.dumps(dict_update) + } + result = self.client_put('/json/user_state', params) + self.assert_json_success(result) + # Assert the data was updated. + updated_dict = initial_dict.copy() + updated_dict.update(dict_update) + result = self.client_get('/json/user_state') + self.assert_json_success(result) + self.assertEqual(result.json()['state'], updated_dict) + # Assert errors on invalid requests. + params = { # type: ignore # Ignore 'incompatible type "str": "List[str]"; expected "str": "str"' for testing + 'keys': ["This is a list, but should be a serialized string."] + } + result = self.client_get('/json/user_state', params) + self.assert_json_error(result, 'Argument "keys" is not valid JSON.') + params = { + 'keys': ujson.dumps(["key 1", "nonexistent key"]) + } + result = self.client_get('/json/user_state', params) + self.assert_json_error(result, "Key does not exist.") + params = { # type: ignore # Ignore 'incompatible type "str": "List[str]"; expected "str": "str"' for testing + 'state': ujson.dumps({'foo': [1, 2, 3]}) + } + result = self.client_put('/json/user_state', params) + self.assert_json_error(result, "Value type is , but should be str.") + # Remove some entries. + keys_to_remove = ['key 1', 'key 2'] + params = { + 'keys': ujson.dumps(keys_to_remove) + } + result = self.client_delete('/json/user_state', params) + self.assert_json_success(result) + # Assert the entries were removed. + for key in keys_to_remove: + updated_dict.pop(key) + result = self.client_get('/json/user_state') + self.assert_json_success(result) + self.assertEqual(result.json()['state'], updated_dict) + # Try to remove an existing and a nonexistent key. + params = { + 'keys': ujson.dumps(['key 3', 'nonexistent key']) + } + result = self.client_delete('/json/user_state', params) + self.assert_json_error(result, "Key does not exist.") + # Assert an error has been thrown and no entries were removed. + result = self.client_get('/json/user_state') + self.assert_json_success(result) + self.assertEqual(result.json()['state'], updated_dict) + # Remove the entire state. + result = self.client_delete('/json/user_state') + self.assert_json_success(result) + # Assert the entire state has been removed. + result = self.client_get('/json/user_state') + self.assert_json_success(result) + self.assertEqual(result.json()['state'], {}) + class TestServiceBotConfigHandler(ZulipTestCase): def setUp(self) -> None: self.user_profile = self.example_user("othello") diff --git a/zerver/views/state.py b/zerver/views/state.py new file mode 100644 index 0000000000..17601e0d6a --- /dev/null +++ b/zerver/views/state.py @@ -0,0 +1,45 @@ +from django.http import HttpRequest, HttpResponse +from django.utils.translation import ugettext as _ +from zerver.lib.bot_storage import ( + get_bot_state, + set_bot_state, + remove_bot_state, + get_keys_in_bot_state, + is_key_in_bot_state, + StateError, +) +from zerver.decorator import has_request_variables, REQ +from zerver.lib.response import json_success, json_error +from zerver.lib.validator import check_dict, check_list, check_string +from zerver.models import UserProfile + +from typing import Dict, List, Optional + +@has_request_variables +def update_state(request, user_profile, state=REQ(validator=check_dict([]))): + # type: (HttpRequest, UserProfile, Optional[Dict[str, str]]) -> HttpResponse + try: + set_bot_state(user_profile, list(state.items())) + except StateError as e: + return json_error(str(e)) + return json_success() + +@has_request_variables +def get_state(request, user_profile, keys=REQ(validator=check_list(check_string), default=None)): + # type: (HttpRequest, UserProfile, Optional[List[str]]) -> HttpResponse + keys = keys or get_keys_in_bot_state(user_profile) + try: + state = {key: get_bot_state(user_profile, key) for key in keys} + except StateError as e: + return json_error(str(e)) + return json_success({'state': state}) + +@has_request_variables +def remove_state(request, user_profile, keys=REQ(validator=check_list(check_string), default=None)): + # type: (HttpRequest, UserProfile, Optional[List[str]]) -> HttpResponse + keys = keys or get_keys_in_bot_state(user_profile) + try: + remove_bot_state(user_profile, keys) + except StateError as e: + return json_error(str(e)) + return json_success() diff --git a/zproject/urls.py b/zproject/urls.py index b9a73248a8..6f310a9916 100644 --- a/zproject/urls.py +++ b/zproject/urls.py @@ -200,6 +200,12 @@ v1_api_and_json_patterns = [ url(r'^user_uploads$', rest_dispatch, {'POST': 'zerver.views.upload.upload_file_backend'}), + # user_state -> zerver.views.state + url(r'^user_state$', rest_dispatch, + {'PUT': 'zerver.views.state.update_state', + 'GET': 'zerver.views.state.get_state', + 'DELETE': 'zerver.views.state.remove_state'}), + # users/me -> zerver.views url(r'^users/me$', rest_dispatch, {'GET': 'zerver.views.users.get_profile_backend',