From 8fa94138bc7899ce2ef40f143cbd4a1bccd4df01 Mon Sep 17 00:00:00 2001
From: Zixuan James Li <p359101898@gmail.com>
Date: Thu, 23 Jun 2022 14:04:16 -0400
Subject: [PATCH] realm_domains: Type QuerySet.values as RealmDomainDict.

This ensure that the return type is compatible with the
actual type of `realm.realmdomain_set.values`.

Signed-off-by: Zixuan James Li <p359101898@gmail.com>
---
 zerver/actions/realm_domains.py | 24 ++++++++++++++----------
 zerver/models.py                |  7 ++++++-
 zerver/tests/test_audit_log.py  | 25 +++++++++++++------------
 3 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/zerver/actions/realm_domains.py b/zerver/actions/realm_domains.py
index 1619d3c448..13778727ec 100644
--- a/zerver/actions/realm_domains.py
+++ b/zerver/actions/realm_domains.py
@@ -9,6 +9,7 @@ from zerver.models import (
     Realm,
     RealmAuditLog,
     RealmDomain,
+    RealmDomainDict,
     UserProfile,
     active_user_ids,
     get_realm_domains,
@@ -23,6 +24,7 @@ def do_add_realm_domain(
     realm_domain = RealmDomain.objects.create(
         realm=realm, domain=domain, allow_subdomains=allow_subdomains
     )
+    added_domain = RealmDomainDict(domain=domain, allow_subdomains=allow_subdomains)
 
     RealmAuditLog.objects.create(
         realm=realm,
@@ -32,7 +34,7 @@ def do_add_realm_domain(
         extra_data=orjson.dumps(
             {
                 "realm_domains": get_realm_domains(realm),
-                "added_domain": {"domain": domain, "allow_subdomains": allow_subdomains},
+                "added_domain": added_domain,
             }
         ).decode(),
     )
@@ -40,7 +42,7 @@ def do_add_realm_domain(
     event = dict(
         type="realm_domains",
         op="add",
-        realm_domain=dict(
+        realm_domain=RealmDomainDict(
             domain=realm_domain.domain, allow_subdomains=realm_domain.allow_subdomains
         ),
     )
@@ -55,6 +57,10 @@ def do_change_realm_domain(
 ) -> None:
     realm_domain.allow_subdomains = allow_subdomains
     realm_domain.save(update_fields=["allow_subdomains"])
+    changed_domain = RealmDomainDict(
+        domain=realm_domain.domain,
+        allow_subdomains=realm_domain.allow_subdomains,
+    )
 
     RealmAuditLog.objects.create(
         realm=realm_domain.realm,
@@ -64,10 +70,7 @@ def do_change_realm_domain(
         extra_data=orjson.dumps(
             {
                 "realm_domains": get_realm_domains(realm_domain.realm),
-                "changed_domain": {
-                    "domain": realm_domain.domain,
-                    "allow_subdomains": realm_domain.allow_subdomains,
-                },
+                "changed_domain": changed_domain,
             }
         ).decode(),
     )
@@ -91,6 +94,10 @@ def do_remove_realm_domain(
     realm = realm_domain.realm
     domain = realm_domain.domain
     realm_domain.delete()
+    removed_domain = RealmDomainDict(
+        domain=realm_domain.domain,
+        allow_subdomains=realm_domain.allow_subdomains,
+    )
 
     RealmAuditLog.objects.create(
         realm=realm,
@@ -100,10 +107,7 @@ def do_remove_realm_domain(
         extra_data=orjson.dumps(
             {
                 "realm_domains": get_realm_domains(realm),
-                "removed_domain": {
-                    "domain": realm_domain.domain,
-                    "allow_subdomains": realm_domain.allow_subdomains,
-                },
+                "removed_domain": removed_domain,
             }
         ).decode(),
     )
diff --git a/zerver/models.py b/zerver/models.py
index f84a8fb575..ce46f65e46 100644
--- a/zerver/models.py
+++ b/zerver/models.py
@@ -1075,7 +1075,12 @@ class EmailContainsPlusError(Exception):
     pass
 
 
-def get_realm_domains(realm: Realm) -> List[Dict[str, Union[str, bool]]]:
+class RealmDomainDict(TypedDict):
+    domain: str
+    allow_subdomains: bool
+
+
+def get_realm_domains(realm: Realm) -> List[RealmDomainDict]:
     return list(realm.realmdomain_set.values("domain", "allow_subdomains"))
 
 
diff --git a/zerver/tests/test_audit_log.py b/zerver/tests/test_audit_log.py
index dff95e63ad..c62b933a98 100644
--- a/zerver/tests/test_audit_log.py
+++ b/zerver/tests/test_audit_log.py
@@ -58,6 +58,7 @@ from zerver.models import (
     Message,
     Realm,
     RealmAuditLog,
+    RealmDomainDict,
     RealmPlayground,
     Recipient,
     Subscription,
@@ -684,10 +685,10 @@ class TestRealmAuditLog(ZulipTestCase):
 
         now = timezone_now()
         realm_domain = do_add_realm_domain(user.realm, "zulip.org", False, acting_user=user)
-        added_domain: Dict[str, Union[str, bool]] = {
-            "domain": "zulip.org",
-            "allow_subdomains": False,
-        }
+        added_domain = RealmDomainDict(
+            domain="zulip.org",
+            allow_subdomains=False,
+        )
         expected_extra_data = {
             "realm_domains": initial_domains + [added_domain],
             "added_domain": added_domain,
@@ -705,10 +706,10 @@ class TestRealmAuditLog(ZulipTestCase):
 
         now = timezone_now()
         do_change_realm_domain(realm_domain, True, acting_user=user)
-        changed_domain: Dict[str, Union[str, bool]] = {
-            "domain": "zulip.org",
-            "allow_subdomains": True,
-        }
+        changed_domain = RealmDomainDict(
+            domain="zulip.org",
+            allow_subdomains=True,
+        )
         expected_extra_data = {
             "realm_domains": initial_domains + [changed_domain],
             "changed_domain": changed_domain,
@@ -726,10 +727,10 @@ class TestRealmAuditLog(ZulipTestCase):
 
         now = timezone_now()
         do_remove_realm_domain(realm_domain, acting_user=user)
-        removed_domain = {
-            "domain": "zulip.org",
-            "allow_subdomains": True,
-        }
+        removed_domain = RealmDomainDict(
+            domain="zulip.org",
+            allow_subdomains=True,
+        )
         expected_extra_data = {
             "realm_domains": initial_domains,
             "removed_domain": removed_domain,