embedded bots: Add views to access state.

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

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 <class 'list'>, 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")

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

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',