diff --git a/tests/rest/admin/test_user.py b/tests/rest/admin/test_user.py
index 553355bc6c..dfc0a0321b 100644
--- a/tests/rest/admin/test_user.py
+++ b/tests/rest/admin/test_user.py
@@ -796,9 +796,9 @@ class UsersListTestCase(unittest.HomeserverTestCase):
 
         # Set avatar URL to all users, that no user has a NULL value to avoid
         # different sort order between SQlite and PostreSQL
-        self.get_success(self.store.set_profile_avatar_url("user1", "mxc://url3"))
-        self.get_success(self.store.set_profile_avatar_url("user2", "mxc://url2"))
-        self.get_success(self.store.set_profile_avatar_url("admin", "mxc://url1"))
+        self.get_success(self.store.set_profile_avatar_url("user1", "mxc://url3", 1))
+        self.get_success(self.store.set_profile_avatar_url("user2", "mxc://url2", 1))
+        self.get_success(self.store.set_profile_avatar_url("admin", "mxc://url1", 1))
 
         # order by default (name)
         self._order_test([self.admin_user, user1, user2], None)
diff --git a/tests/rest/client/test_room_access_rules.py b/tests/rest/client/test_room_access_rules.py
index 64744b7176..f3ebfa3c2f 100644
--- a/tests/rest/client/test_room_access_rules.py
+++ b/tests/rest/client/test_room_access_rules.py
@@ -23,7 +23,7 @@ from twisted.internet import defer
 from synapse.api.constants import EventTypes, JoinRules, Membership, RoomCreationPreset
 from synapse.api.errors import SynapseError
 from synapse.rest import admin
-from synapse.rest.client.v1 import directory, login, room
+from synapse.rest.client import directory, login, room
 from synapse.third_party_rules.access_rules import (
     ACCESS_RULES_TYPE,
     FROZEN_STATE_TYPE,
diff --git a/tests/rest/client/v2_alpha/test_register.py b/tests/rest/client/v2_alpha/test_register.py
index d8edbe0ac2..ec1009abc5 100644
--- a/tests/rest/client/v2_alpha/test_register.py
+++ b/tests/rest/client/v2_alpha/test_register.py
@@ -28,7 +28,17 @@ import synapse.rest.admin
 from synapse.api.constants import APP_SERVICE_REGISTRATION_TYPE, LoginType
 from synapse.api.errors import Codes
 from synapse.appservice import ApplicationService
-from synapse.rest.client import account, account_validity, login, logout, register, sync
+from synapse.rest.client import (
+    account,
+    account_validity,
+    login,
+    logout,
+    profile,
+    register,
+    room,
+    sync,
+    user_directory,
+)
 
 from tests import unittest
 from tests.unittest import override_config
@@ -572,9 +582,9 @@ class AccountValidityTestCase(unittest.HomeserverTestCase):
 class AccountValidityUserDirectoryTestCase(unittest.HomeserverTestCase):
 
     servlets = [
-        synapse.rest.client.v1.profile.register_servlets,
-        synapse.rest.client.v1.room.register_servlets,
-        synapse.rest.client.v2_alpha.user_directory.register_servlets,
+        profile.register_servlets,
+        room.register_servlets,
+        user_directory.register_servlets,
         login.register_servlets,
         register.register_servlets,
         synapse.rest.admin.register_servlets_for_client_rest_resource,
diff --git a/tests/rulecheck/test_domainrulecheck.py b/tests/rulecheck/test_domainrulecheck.py
index 3c70a9c8c0..29e919f087 100644
--- a/tests/rulecheck/test_domainrulecheck.py
+++ b/tests/rulecheck/test_domainrulecheck.py
@@ -18,7 +18,7 @@ import json
 
 import synapse.rest.admin
 from synapse.config._base import ConfigError
-from synapse.rest.client.v1 import login, room
+from synapse.rest.client import login, room
 from synapse.rulecheck.domain_rule_checker import DomainRuleChecker
 
 from tests import unittest
diff --git a/tests/storage/test_profile.py b/tests/storage/test_profile.py
index ae29357b0b..33ef510b21 100644
--- a/tests/storage/test_profile.py
+++ b/tests/storage/test_profile.py
@@ -27,7 +27,7 @@ class ProfileStoreTestCase(unittest.HomeserverTestCase):
         self.get_success(self.store.create_profile(self.u_frank.localpart))
 
         self.get_success(
-            self.store.set_profile_displayname(self.u_frank.localpart, "Frank")
+            self.store.set_profile_displayname(self.u_frank.localpart, "Frank", 1)
         )
 
         self.assertEquals(
@@ -68,7 +68,7 @@ class ProfileStoreTestCase(unittest.HomeserverTestCase):
 
         # test set to None
         self.get_success(
-            self.store.set_profile_avatar_url(self.u_frank.localpart, None)
+            self.store.set_profile_avatar_url(self.u_frank.localpart, None, 1)
         )
 
         self.assertIsNone(