update tests to use enum

5 files changed, 33 insertions, 15 deletions

diff --git a/synapse/storage/databases/main/experimental_features.py b/synapse/storage/databases/main/experimental_features.py
index f88598794a..e751f68600 100644
--- a/synapse/storage/databases/main/experimental_features.py
+++ b/synapse/storage/databases/main/experimental_features.py
@@ -104,4 +104,4 @@ class ExperimentalFeaturesStore(CacheInvalidationWorkerStore):
             res["enabled"] = True
 
         return res["enabled"]
->>>>>>> 6c2267f83d... add a db function to tell if just one feature is enabled
+
diff --git a/tests/handlers/test_presence.py b/tests/handlers/test_presence.py
index 7a71db3af6..ab763d6284 100644
--- a/tests/handlers/test_presence.py
+++ b/tests/handlers/test_presence.py
@@ -36,6 +36,7 @@ from synapse.handlers.presence import (
     handle_update,
 )
 from synapse.rest import admin
+from synapse.rest.admin.experimental_features import ExperimentalFeature
 from synapse.rest.client import room
 from synapse.server import HomeServer
 from synapse.types import JsonDict, UserID, get_domain_from_id
@@ -514,9 +515,14 @@ class PresenceTimeoutTestCase(unittest.TestCase):
 
 
 class PresenceHandlerTestCase(BaseMultiWorkerStreamTestCase):
+    servlets = [
+        admin.register_servlets,
+    ]
+
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.presence_handler = hs.get_presence_handler()
         self.clock = hs.get_clock()
+        self.user = self.register_user("test", "pass")
 
     def test_external_process_timeout(self) -> None:
         """Test that if an external process doesn't update the records for a while
@@ -734,13 +740,12 @@ class PresenceHandlerTestCase(BaseMultiWorkerStreamTestCase):
             test_with_workers: If True, check the presence state of the user by calling
                 /sync against a worker, rather than the main process.
         """
-        user_id = "@test:server"
         status_msg = "I'm busy!"
 
         # set busy state in db
         self.get_success(
-            self.hs.get_datastores().main.set_feature_for_user(
-                user_id, "msc_3026", True
+            self.hs.get_datastores().main.set_features_for_user(
+                self.user, {ExperimentalFeature.MSC3026: True}
             )
         )
 
@@ -755,20 +760,22 @@ class PresenceHandlerTestCase(BaseMultiWorkerStreamTestCase):
             )
 
         # Set presence to BUSY
-        self._set_presencestate_with_status_msg(user_id, PresenceState.BUSY, status_msg)
+        self._set_presencestate_with_status_msg(
+            self.user, PresenceState.BUSY, status_msg
+        )
 
         # Perform a sync with a presence state other than busy. This should NOT change
         # our presence status; we only change from busy if we explicitly set it via
         # /presence/*.
         self.get_success(
             worker_to_sync_against.get_presence_handler().user_syncing(
-                user_id, True, PresenceState.ONLINE
+                self.user, True, PresenceState.ONLINE
             )
         )
 
         # Check against the main process that the user's presence did not change.
         state = self.get_success(
-            self.presence_handler.get_state(UserID.from_string(user_id))
+            self.presence_handler.get_state(UserID.from_string(self.user))
         )
         # we should still be busy
         self.assertEqual(state.state, PresenceState.BUSY)
diff --git a/tests/push/test_http.py b/tests/push/test_http.py
index 67c1934eaa..97bc48f246 100644
--- a/tests/push/test_http.py
+++ b/tests/push/test_http.py
@@ -20,6 +20,7 @@ from twisted.test.proto_helpers import MemoryReactor
 import synapse.rest.admin
 from synapse.logging.context import make_deferred_yieldable
 from synapse.push import PusherConfig, PusherConfigException
+from synapse.rest.admin.experimental_features import ExperimentalFeature
 from synapse.rest.client import login, push_rule, pusher, receipts, room
 from synapse.server import HomeServer
 from synapse.types import JsonDict
@@ -829,7 +830,9 @@ class HTTPPusherTests(HomeserverTestCase):
 
         # enable msc3881 per_user flag
         self.get_success(
-            self.hs.get_datastores().main.set_feature_for_user(user_id, "msc3881", True)
+            self.hs.get_datastores().main.set_features_for_user(
+                user_id, {ExperimentalFeature.MSC3881: True}
+            )
         )
 
         # Send a message and check that it generated a push.
@@ -863,7 +866,9 @@ class HTTPPusherTests(HomeserverTestCase):
 
         # enable msc3881 per_user flag
         self.get_success(
-            self.hs.get_datastores().main.set_feature_for_user(user_id, "msc3881", True)
+            self.hs.get_datastores().main.set_features_for_user(
+                user_id, {ExperimentalFeature.MSC3881: True}
+            )
         )
 
         # Send a message and check that it did not generate a push.
@@ -897,7 +902,9 @@ class HTTPPusherTests(HomeserverTestCase):
 
         # enable msc3881 per_user flag
         self.get_success(
-            self.hs.get_datastores().main.set_feature_for_user(user_id, "msc3881", True)
+            self.hs.get_datastores().main.set_features_for_user(
+                user_id, {ExperimentalFeature.MSC3881: True}
+            )
         )
 
         channel = self.make_request("GET", "/pushers", access_token=access_token)
@@ -944,7 +951,9 @@ class HTTPPusherTests(HomeserverTestCase):
 
         # enable msc3881 per_user flag
         self.get_success(
-            self.hs.get_datastores().main.set_feature_for_user(user, "msc3881", True)
+            self.hs.get_datastores().main.set_features_for_user(
+                user, {ExperimentalFeature.MSC3881: True}
+            )
         )
 
         # We create the pusher with an HTTP request rather than with
diff --git a/tests/rest/client/test_keys.py b/tests/rest/client/test_keys.py
index 713c7abec3..ef2800ee54 100644
--- a/tests/rest/client/test_keys.py
+++ b/tests/rest/client/test_keys.py
@@ -23,6 +23,7 @@ from signedjson.sign import sign_json
 
 from synapse.api.errors import Codes
 from synapse.rest import admin
+from synapse.rest.admin.experimental_features import ExperimentalFeature
 from synapse.rest.client import keys, login
 from synapse.types import JsonDict
 
@@ -217,8 +218,8 @@ class KeyQueryTestCase(unittest.HomeserverTestCase):
 
         # enable msc3967 in db
         self.get_success(
-            self.hs.get_datastores().main.set_feature_for_user(
-                alice_id, "msc3967", True
+            self.hs.get_datastores().main.set_features_for_user(
+                alice_id, {ExperimentalFeature.MSC3967: True}
             )
         )
 
diff --git a/tests/rest/client/test_sync.py b/tests/rest/client/test_sync.py
index 773765f789..49cd4cee60 100644
--- a/tests/rest/client/test_sync.py
+++ b/tests/rest/client/test_sync.py
@@ -28,6 +28,7 @@ from synapse.api.constants import (
     ReceiptTypes,
     RelationTypes,
 )
+from synapse.rest.admin.experimental_features import ExperimentalFeature
 from synapse.rest.client import devices, knock, login, read_marker, receipts, room, sync
 from synapse.server import HomeServer
 from synapse.types import JsonDict
@@ -588,8 +589,8 @@ class UnreadMessagesTestCase(unittest.HomeserverTestCase):
         """Tests that /sync returns the right value for the unread count (MSC2654)."""
         # add per-user flag to the DB
         self.get_success(
-            self.hs.get_datastores().main.set_feature_for_user(
-                self.user_id, "msc2654", True
+            self.hs.get_datastores().main.set_features_for_user(
+                self.user_id, {ExperimentalFeature.MSC2654: True}
             )
         )