1 files changed, 17 insertions, 44 deletions

diff --git a/tests/storage/test_room.py b/tests/storage/test_room.py
index bc8400f240..0089d33c93 100644
--- a/tests/storage/test_room.py
+++ b/tests/storage/test_room.py
@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-# Copyright 2014-2016 OpenMarket Ltd
+# Copyright 2014-2021 The Matrix.org Foundation C.I.C.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,22 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-from twisted.internet import defer
-
 from synapse.api.constants import EventTypes
 from synapse.api.room_versions import RoomVersions
 from synapse.types import RoomAlias, RoomID, UserID
 
-from tests import unittest
-from tests.utils import setup_test_homeserver
-
+from tests.unittest import HomeserverTestCase
 
-class RoomStoreTestCase(unittest.TestCase):
-    @defer.inlineCallbacks
-    def setUp(self):
-        hs = yield setup_test_homeserver(self.addCleanup)
 
+class RoomStoreTestCase(HomeserverTestCase):
+    def prepare(self, reactor, clock, hs):
         # We can't test RoomStore on its own without the DirectoryStore, for
         # management of the 'room_aliases' table
         self.store = hs.get_datastore()
@@ -37,7 +30,7 @@ class RoomStoreTestCase(unittest.TestCase):
         self.alias = RoomAlias.from_string("#a-room-name:test")
         self.u_creator = UserID.from_string("@creator:test")
 
-        yield defer.ensureDeferred(
+        self.get_success(
             self.store.store_room(
                 self.room.to_string(),
                 room_creator_user_id=self.u_creator.to_string(),
@@ -46,7 +39,6 @@ class RoomStoreTestCase(unittest.TestCase):
             )
         )
 
-    @defer.inlineCallbacks
     def test_get_room(self):
         self.assertDictContainsSubset(
             {
@@ -54,16 +46,12 @@ class RoomStoreTestCase(unittest.TestCase):
                 "creator": self.u_creator.to_string(),
                 "is_public": True,
             },
-            (yield defer.ensureDeferred(self.store.get_room(self.room.to_string()))),
+            (self.get_success(self.store.get_room(self.room.to_string()))),
         )
 
-    @defer.inlineCallbacks
     def test_get_room_unknown_room(self):
-        self.assertIsNone(
-            (yield defer.ensureDeferred(self.store.get_room("!uknown:test")))
-        )
+        self.assertIsNone((self.get_success(self.store.get_room("!uknown:test"))))
 
-    @defer.inlineCallbacks
     def test_get_room_with_stats(self):
         self.assertDictContainsSubset(
             {
@@ -71,29 +59,17 @@ class RoomStoreTestCase(unittest.TestCase):
                 "creator": self.u_creator.to_string(),
                 "public": True,
             },
-            (
-                yield defer.ensureDeferred(
-                    self.store.get_room_with_stats(self.room.to_string())
-                )
-            ),
+            (self.get_success(self.store.get_room_with_stats(self.room.to_string()))),
         )
 
-    @defer.inlineCallbacks
     def test_get_room_with_stats_unknown_room(self):
         self.assertIsNone(
-            (
-                yield defer.ensureDeferred(
-                    self.store.get_room_with_stats("!uknown:test")
-                )
-            ),
+            (self.get_success(self.store.get_room_with_stats("!uknown:test"))),
         )
 
 
-class RoomEventsStoreTestCase(unittest.TestCase):
-    @defer.inlineCallbacks
-    def setUp(self):
-        hs = setup_test_homeserver(self.addCleanup)
-
+class RoomEventsStoreTestCase(HomeserverTestCase):
+    def prepare(self, reactor, clock, hs):
         # Room events need the full datastore, for persist_event() and
         # get_room_state()
         self.store = hs.get_datastore()
@@ -102,7 +78,7 @@ class RoomEventsStoreTestCase(unittest.TestCase):
 
         self.room = RoomID.from_string("!abcde:test")
 
-        yield defer.ensureDeferred(
+        self.get_success(
             self.store.store_room(
                 self.room.to_string(),
                 room_creator_user_id="@creator:text",
@@ -111,23 +87,21 @@ class RoomEventsStoreTestCase(unittest.TestCase):
             )
         )
 
-    @defer.inlineCallbacks
     def inject_room_event(self, **kwargs):
-        yield defer.ensureDeferred(
+        self.get_success(
             self.storage.persistence.persist_event(
                 self.event_factory.create_event(room_id=self.room.to_string(), **kwargs)
             )
         )
 
-    @defer.inlineCallbacks
     def STALE_test_room_name(self):
         name = "A-Room-Name"
 
-        yield self.inject_room_event(
+        self.inject_room_event(
             etype=EventTypes.Name, name=name, content={"name": name}, depth=1
         )
 
-        state = yield defer.ensureDeferred(
+        state = self.get_success(
             self.store.get_current_state(room_id=self.room.to_string())
         )
 
@@ -137,15 +111,14 @@ class RoomEventsStoreTestCase(unittest.TestCase):
             state[0],
         )
 
-    @defer.inlineCallbacks
     def STALE_test_room_topic(self):
         topic = "A place for things"
 
-        yield self.inject_room_event(
+        self.inject_room_event(
             etype=EventTypes.Topic, topic=topic, content={"topic": topic}, depth=1
         )
 
-        state = yield defer.ensureDeferred(
+        state = self.get_success(
             self.store.get_current_state(room_id=self.room.to_string())
         )