4 files changed, 5 insertions, 3 deletions
diff --git a/tests/storage/test_appservice.py b/tests/storage/test_appservice.py
index 622b16a071..dfeea24599 100644
--- a/tests/storage/test_appservice.py
+++ b/tests/storage/test_appservice.py
@@ -24,7 +24,7 @@ from twisted.internet import defer
 
 from synapse.appservice import ApplicationService, ApplicationServiceState
 from synapse.config._base import ConfigError
-from synapse.storage.appservice import (
+from synapse.storage.data_stores.main.appservice import (
     ApplicationServiceStore,
     ApplicationServiceTransactionStore,
 )
diff --git a/tests/storage/test_cleanup_extrems.py b/tests/storage/test_cleanup_extrems.py
index 34f9c72709..69dcaa63d5 100644
--- a/tests/storage/test_cleanup_extrems.py
+++ b/tests/storage/test_cleanup_extrems.py
@@ -50,6 +50,8 @@ class CleanupExtremBackgroundUpdateStoreTestCase(HomeserverTestCase):
 
         schema_path = os.path.join(
             prepare_database.dir_path,
+            "data_stores",
+            "main",
             "schema",
             "delta",
             "54",
diff --git a/tests/storage/test_profile.py b/tests/storage/test_profile.py
index 45824bd3b2..24c7fe16c3 100644
--- a/tests/storage/test_profile.py
+++ b/tests/storage/test_profile.py
@@ -16,7 +16,7 @@
 
 from twisted.internet import defer
 
-from synapse.storage.profile import ProfileStore
+from synapse.storage.data_stores.main.profile import ProfileStore
 from synapse.types import UserID
 
 from tests import unittest
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index d7d244ce97..7eea57c0e2 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -15,7 +15,7 @@
 
 from twisted.internet import defer
 
-from synapse.storage import UserDirectoryStore
+from synapse.storage.data_stores.main.user_directory import UserDirectoryStore
 
 from tests import unittest
 from tests.utils import setup_test_homeserver