1 files changed, 12 insertions, 15 deletions
diff --git a/tests/handlers/test_directory.py b/tests/handlers/test_directory.py
index 72a2b1443a..58ecf561ff 100644
--- a/tests/handlers/test_directory.py
+++ b/tests/handlers/test_directory.py
@@ -25,6 +25,8 @@ from synapse.http.client import HttpClient
 from synapse.handlers.directory import DirectoryHandler
 from synapse.storage.directory import RoomAliasMapping
 
+from tests.utils import SQLiteMemoryDbPool
+
 
 logging.getLogger().addHandler(logging.NullHandler())
 
@@ -37,6 +39,7 @@ class DirectoryHandlers(object):
 class DirectoryTestCase(unittest.TestCase):
     """ Tests the directory service. """
 
+    @defer.inlineCallbacks
     def setUp(self):
         self.mock_federation = Mock(spec=[
             "make_query",
@@ -48,32 +51,27 @@ class DirectoryTestCase(unittest.TestCase):
         self.mock_federation.register_query_handler = register_query_handler
 
         hs = HomeServer("test",
-            datastore=Mock(spec=[
-                "get_association_from_room_alias",
-                "get_joined_hosts_for_room",
-            ]),
+            db_pool=SQLiteMemoryDbPool(),
             http_client=None,
             resource_for_federation=Mock(),
             replication_layer=self.mock_federation,
         )
         hs.handlers = DirectoryHandlers(hs)
 
-        self.handler = hs.get_handlers().directory_handler
+        yield hs.get_db_pool().prepare()
 
-        self.datastore = hs.get_datastore()
+        self.handler = hs.get_handlers().directory_handler
 
-        def hosts(room_id):
-            return defer.succeed([])
-        self.datastore.get_joined_hosts_for_room.side_effect = hosts
+        self.store = hs.get_datastore()
 
         self.my_room = hs.parse_roomalias("#my-room:test")
+        self.your_room = hs.parse_roomalias("#your-room:test")
         self.remote_room = hs.parse_roomalias("#another:remote")
 
     @defer.inlineCallbacks
     def test_get_local_association(self):
-        mocked_get = self.datastore.get_association_from_room_alias
-        mocked_get.return_value = defer.succeed(
-            RoomAliasMapping("!8765qwer:test", "#my-room:test", ["test"])
+        yield self.store.create_room_alias_association(
+            self.my_room, "!8765qwer:test", ["test"]
         )
 
         result = yield self.handler.get_association(self.my_room)
@@ -106,9 +104,8 @@ class DirectoryTestCase(unittest.TestCase):
 
     @defer.inlineCallbacks
     def test_incoming_fed_query(self):
-        mocked_get = self.datastore.get_association_from_room_alias
-        mocked_get.return_value = defer.succeed(
-            RoomAliasMapping("!8765asdf:test", "#your-room:test", ["test"])
+        yield self.store.create_room_alias_association(
+            self.your_room, "!8765asdf:test", ["test"]
         )
 
         response = yield self.query_handlers["directory"](