3 files changed, 12 insertions, 5 deletions
diff --git a/tests/storage/test_appservice.py b/tests/storage/test_appservice.py
index 5abecdf6e0..ed8af10d87 100644
--- a/tests/storage/test_appservice.py
+++ b/tests/storage/test_appservice.py
@@ -439,7 +439,7 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
         f2 = self._write_config(suffix="2")
 
         config = Mock(app_service_config_files=[f1, f2])
-        hs = yield setup_test_homeserver(config=config)
+        hs = yield setup_test_homeserver(config=config, datastore=Mock())
 
         ApplicationServiceStore(hs)
 
@@ -449,7 +449,7 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
         f2 = self._write_config(id="id", suffix="2")
 
         config = Mock(app_service_config_files=[f1, f2])
-        hs = yield setup_test_homeserver(config=config)
+        hs = yield setup_test_homeserver(config=config, datastore=Mock())
 
         with self.assertRaises(ConfigError) as cm:
             ApplicationServiceStore(hs)
@@ -465,7 +465,7 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
         f2 = self._write_config(as_token="as_token", suffix="2")
 
         config = Mock(app_service_config_files=[f1, f2])
-        hs = yield setup_test_homeserver(config=config)
+        hs = yield setup_test_homeserver(config=config, datastore=Mock())
 
         with self.assertRaises(ConfigError) as cm:
             ApplicationServiceStore(hs)
diff --git a/tests/storage/test_registration.py b/tests/storage/test_registration.py
index a35efcc71e..7b3b4c13bc 100644
--- a/tests/storage/test_registration.py
+++ b/tests/storage/test_registration.py
@@ -18,7 +18,6 @@ from tests import unittest
 from twisted.internet import defer
 
 from synapse.api.errors import StoreError
-from synapse.storage.registration import RegistrationStore
 from synapse.util import stringutils
 
 from tests.utils import setup_test_homeserver
@@ -31,7 +30,7 @@ class RegistrationStoreTestCase(unittest.TestCase):
         hs = yield setup_test_homeserver()
         self.db_pool = hs.get_db_pool()
 
-        self.store = RegistrationStore(hs)
+        self.store = hs.get_datastore()
 
         self.user_id = "@my-user:test"
         self.tokens = ["AbCdEfGhIjKlMnOpQrStUvWxYz",
diff --git a/tests/utils.py b/tests/utils.py
index d75d492cb5..43cc2b30cd 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -60,8 +60,10 @@ def setup_test_homeserver(name="test", datastore=None, config=None, **kargs):
             name, db_pool=db_pool, config=config,
             version_string="Synapse/tests",
             database_engine=create_engine("sqlite3"),
+            get_db_conn=db_pool.get_db_conn,
             **kargs
         )
+        hs.setup()
     else:
         hs = HomeServer(
             name, db_pool=None, datastore=datastore, config=config,
@@ -280,6 +282,12 @@ class SQLiteMemoryDbPool(ConnectionPool, object):
             lambda conn: prepare_database(conn, engine)
         )
 
+    def get_db_conn(self):
+        conn = self.connect()
+        engine = create_engine("sqlite3")
+        prepare_database(conn, engine)
+        return conn
+
 
 class MemoryDataStore(object):