1 files changed, 4 insertions, 3 deletions
diff --git a/tests/handlers/test_directory.py b/tests/handlers/test_directory.py
index 58ecf561ff..f7eb7b4f03 100644
--- a/tests/handlers/test_directory.py
+++ b/tests/handlers/test_directory.py
@@ -50,16 +50,17 @@ class DirectoryTestCase(unittest.TestCase):
             self.query_handlers[query_type] = handler
         self.mock_federation.register_query_handler = register_query_handler
 
+        db_pool = SQLiteMemoryDbPool()
+        yield db_pool.prepare()
+
         hs = HomeServer("test",
-            db_pool=SQLiteMemoryDbPool(),
+            db_pool=db_pool,
             http_client=None,
             resource_for_federation=Mock(),
             replication_layer=self.mock_federation,
         )
         hs.handlers = DirectoryHandlers(hs)
 
-        yield hs.get_db_pool().prepare()
-
         self.handler = hs.get_handlers().directory_handler
 
         self.store = hs.get_datastore()