From f03605bd737bf6812cf24f58ccde2adc65f84a12 Mon Sep 17 00:00:00 2001
From: Steve Howell <showell@zulip.com>
Date: Thu, 23 Jul 2020 16:01:27 +0000
Subject: [PATCH] event_schema: Support plan_type in check_realm_update.

---
 zerver/lib/event_schema.py  | 20 +++++++++++++++++++-
 zerver/tests/test_events.py | 11 +----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/zerver/lib/event_schema.py b/zerver/lib/event_schema.py
index 029b81fbbd..04c849f9e7 100644
--- a/zerver/lib/event_schema.py
+++ b/zerver/lib/event_schema.py
@@ -349,6 +349,13 @@ def check_realm_bot_update(var_name: str, event: Dict[str, Any], field: str,) ->
     assert {"user_id", field} == set(event["bot"].keys())
 
 
+_check_plan_type_extra_data = check_dict_only(
+    required_keys=[
+        # force vertical
+        ("upload_quota", check_int),
+    ]
+)
+
 """
 realm/update events are flexible for values;
 we will use a more strict checker to check
@@ -360,7 +367,11 @@ _check_realm_update = check_events_dict(
         ("op", equals("update")),
         ("property", check_string),
         ("value", check_value),
-    ]
+    ],
+    optional_keys=[
+        # force vertical
+        ("extra_data", _check_plan_type_extra_data),
+    ],
 )
 
 
@@ -381,6 +392,13 @@ def check_realm_update(var_name: str, event: Dict[str, Any], prop: str,) -> None
     assert prop == event["property"]
     value = event["value"]
 
+    if prop == "plan_type":
+        assert isinstance(value, int)
+        assert "extra_data" in event.keys()
+        return
+
+    assert "extra_data" not in event.keys()
+
     if prop in ["notifications_stream_id", "signup_notifications_stream_id"]:
         assert isinstance(value, int)
         return
diff --git a/zerver/tests/test_events.py b/zerver/tests/test_events.py
index b3b8fd77a1..11a2c33647 100644
--- a/zerver/tests/test_events.py
+++ b/zerver/tests/test_events.py
@@ -1326,18 +1326,9 @@ class NormalActionsTest(BaseAction):
         self.assertEqual(state_data['realm_plan_type'], Realm.SELF_HOSTED)
         self.assertEqual(state_data['zulip_plan_is_not_limited'], True)
 
-        schema_checker = check_events_dict([
-            ('type', equals('realm')),
-            ('op', equals('update')),
-            ('property', equals('plan_type')),
-            ('value', equals(Realm.LIMITED)),
-            ('extra_data', check_dict_only([
-                ('upload_quota', check_int),
-            ])),
-        ])
         events = self.verify_action(
             lambda: do_change_plan_type(realm, Realm.LIMITED))
-        schema_checker('events[0]', events[0])
+        check_realm_update('events[0]', events[0], 'plan_type')
 
         state_data = fetch_initial_state_data(self.user_profile, None, "", False, False)
         self.assertEqual(state_data['realm_plan_type'], Realm.LIMITED)