1 files changed, 2 insertions, 2 deletions
diff --git a/tests/storage/test_keys.py b/tests/storage/test_keys.py
index a94b5fd721..9059095525 100644
--- a/tests/storage/test_keys.py
+++ b/tests/storage/test_keys.py
@@ -37,7 +37,7 @@ KEY_2 = decode_verify_key_base64(
 
 class KeyStoreTestCase(tests.unittest.HomeserverTestCase):
     def test_get_server_verify_keys(self):
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
 
         key_id_1 = "ed25519:key1"
         key_id_2 = "ed25519:KEY_ID_2"
@@ -74,7 +74,7 @@ class KeyStoreTestCase(tests.unittest.HomeserverTestCase):
     def test_cache(self):
         """Check that updates correctly invalidate the cache."""
 
-        store = self.hs.get_datastore()
+        store = self.hs.get_datastores().main
 
         key_id_1 = "ed25519:key1"
         key_id_2 = "ed25519:key2"