Add type annotations to tests.storage.test_appservice. (#11488)

matrix-org · Dec 2, 2021 · 435f044 · 435f044
1 parent f61462e
commit 435f044
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 53 deletions.
diff --git a/changelog.d/11488.misc b/changelog.d/11488.misc
@@ -0,0 +1 @@
+Add type annotations to `tests.storage.test_appservice`.
diff --git a/mypy.ini b/mypy.ini
@@ -111,7 +111,6 @@ exclude = (?x)
    |tests/server_notices/test_resource_limits_server_notices.py
    |tests/state/test_v2.py
    |tests/storage/test_account_data.py
-   |tests/storage/test_appservice.py
    |tests/storage/test_background_update.py
    |tests/storage/test_base.py
    |tests/storage/test_client_ips.py

diff --git a/synapse/appservice/__init__.py b/synapse/appservice/__init__.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 import logging
 import re
+from enum import Enum
 from typing import TYPE_CHECKING, Iterable, List, Match, Optional
 
 from synapse.api.constants import EventTypes
@@ -27,7 +28,7 @@
 logger = logging.getLogger(__name__)
 
 
-class ApplicationServiceState:
+class ApplicationServiceState(Enum):
     DOWN = "down"
     UP = "up"
 

diff --git a/synapse/storage/databases/main/appservice.py b/synapse/storage/databases/main/appservice.py
@@ -143,7 +143,7 @@ async def get_appservices_by_state(
             A list of ApplicationServices, which may be empty.
         """
         results = await self.db_pool.simple_select_list(
-            "application_services_state", {"state": state}, ["as_id"]
+            "application_services_state", {"state": state.value}, ["as_id"]
         )
         # NB: This assumes this class is linked with ApplicationServiceStore
         as_list = self.get_app_services()
@@ -173,7 +173,7 @@ async def get_appservice_state(
             desc="get_appservice_state",
         )
         if result:
-            return result.get("state")
+            return ApplicationServiceState(result.get("state"))
         return None
 
     async def set_appservice_state(
@@ -186,7 +186,7 @@ async def set_appservice_state(
             state: The connectivity state to apply.
         """
         await self.db_pool.simple_upsert(
-            "application_services_state", {"as_id": service.id}, {"state": state}
+            "application_services_state", {"as_id": service.id}, {"state": state.value}
         )
 
     async def create_appservice_txn(