embedded bots: Add StateHandler remove() function.

zerver/lib/bot_lib.py

@@ -8,7 +8,7 @@ import re
 import importlib
 from zerver.lib.actions import internal_send_message
 from zerver.models import UserProfile, \
-    get_bot_state, set_bot_state, get_bot_state_size, is_key_in_bot_state
+    get_bot_state, set_bot_state, get_bot_state_size, is_key_in_bot_state, remove_bot_state
 from zerver.lib.integrations import EMBEDDED_BOTS
 
 from six.moves import configparser
@@ -67,6 +67,10 @@ class StateHandler(object):
                                         "should be str.".format(type(marshaled_value)))
             set_bot_state(self.user_profile, key, marshaled_value)
 
+    def remove(self, key):
+        # type: (Text) -> None
+        remove_bot_state(self.user_profile, key)
+
     def contains(self, key):
         # type: (Text) -> bool
         return is_key_in_bot_state(self.user_profile, key)

zerver/models.py

@@ -1954,6 +1954,10 @@ def set_bot_state(bot_profile, key, value):
         obj.value = value
         obj.save()
 
+def remove_bot_state(bot_profile, key):
+    # type: (UserProfile, Text) -> None
+    removed_ctr, removed_entries = BotUserStateData.objects.get(bot_profile=bot_profile, key=key).delete()
+
 def is_key_in_bot_state(bot_profile, key):
     # type: (UserProfile, Text) -> bool
     return BotUserStateData.objects.filter(bot_profile=bot_profile, key=key).exists()

zerver/tests/test_service_bot_system.py

@@ -193,6 +193,18 @@ class TestServiceBotStateHandler(ZulipTestCase):
         second_storage.put('another big entry', 'x' * (StateHandler.state_size_limit - 40))
         second_storage.put('normal entry', 'abcd')
 
+    def test_entry_removal(self):
+        # type: () -> None
+        storage = StateHandler(self.bot_profile)
+        storage.put('some key', 'some value')
+        storage.put('another key', 'some value')
+        self.assertTrue(storage.contains('some key'))
+        self.assertTrue(storage.contains('another key'))
+        storage.remove('some key')
+        self.assertFalse(storage.contains('some key'))
+        self.assertTrue(storage.contains('another key'))
+        self.assertRaises(BotUserStateData.DoesNotExist, lambda: storage.remove('some key'))
+
 class TestServiceBotEventTriggers(ZulipTestCase):
 
     def setUp(self):