summary refs log tree commit diff
path: root/tests/storage
diff options
context:
space:
mode:
Diffstat (limited to 'tests/storage')
-rw-r--r--tests/storage/test_appservice.py131
-rw-r--r--tests/storage/test_base.py4
-rw-r--r--tests/storage/test_client_ips.py160
-rw-r--r--tests/storage/test_directory.py3
-rw-r--r--tests/storage/test_event_federation.py6
-rw-r--r--tests/storage/test_monthly_active_users.py211
-rw-r--r--tests/storage/test_presence.py7
-rw-r--r--tests/storage/test_profile.py2
-rw-r--r--tests/storage/test_purge.py106
-rw-r--r--tests/storage/test_state.py141
-rw-r--r--tests/storage/test_transactions.py45
-rw-r--r--tests/storage/test_user_directory.py2
12 files changed, 578 insertions, 240 deletions
diff --git a/tests/storage/test_appservice.py b/tests/storage/test_appservice.py
index c893990454..3f0083831b 100644
--- a/tests/storage/test_appservice.py
+++ b/tests/storage/test_appservice.py
@@ -37,18 +37,14 @@ class ApplicationServiceStoreTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def setUp(self):
         self.as_yaml_files = []
-        config = Mock(
-            app_service_config_files=self.as_yaml_files,
-            event_cache_size=1,
-            password_providers=[],
-        )
         hs = yield setup_test_homeserver(
-            self.addCleanup,
-            config=config,
-            federation_sender=Mock(),
-            federation_client=Mock(),
+            self.addCleanup, federation_sender=Mock(), federation_client=Mock()
         )
 
+        hs.config.app_service_config_files = self.as_yaml_files
+        hs.config.event_cache_size = 1
+        hs.config.password_providers = []
+
         self.as_token = "token1"
         self.as_url = "some_url"
         self.as_id = "as1"
@@ -58,7 +54,7 @@ class ApplicationServiceStoreTestCase(unittest.TestCase):
         self._add_appservice("token2", "as2", "some_url", "some_hs_token", "bob")
         self._add_appservice("token3", "as3", "some_url", "some_hs_token", "bob")
         # must be done after inserts
-        self.store = ApplicationServiceStore(None, hs)
+        self.store = ApplicationServiceStore(hs.get_db_conn(), hs)
 
     def tearDown(self):
         # TODO: suboptimal that we need to create files for tests!
@@ -105,18 +101,16 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
     def setUp(self):
         self.as_yaml_files = []
 
-        config = Mock(
-            app_service_config_files=self.as_yaml_files,
-            event_cache_size=1,
-            password_providers=[],
-        )
         hs = yield setup_test_homeserver(
-            self.addCleanup,
-            config=config,
-            federation_sender=Mock(),
-            federation_client=Mock(),
+            self.addCleanup, federation_sender=Mock(), federation_client=Mock()
         )
+
+        hs.config.app_service_config_files = self.as_yaml_files
+        hs.config.event_cache_size = 1
+        hs.config.password_providers = []
+
         self.db_pool = hs.get_db_pool()
+        self.engine = hs.database_engine
 
         self.as_list = [
             {"token": "token1", "url": "https://matrix-as.org", "id": "id_1"},
@@ -129,7 +123,7 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
 
         self.as_yaml_files = []
 
-        self.store = TestTransactionStore(None, hs)
+        self.store = TestTransactionStore(hs.get_db_conn(), hs)
 
     def _add_service(self, url, as_token, id):
         as_yaml = dict(
@@ -146,29 +140,35 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
             self.as_yaml_files.append(as_token)
 
     def _set_state(self, id, state, txn=None):
-        return self.db_pool.runQuery(
-            "INSERT INTO application_services_state(as_id, state, last_txn) "
-            "VALUES(?,?,?)",
+        return self.db_pool.runOperation(
+            self.engine.convert_param_style(
+                "INSERT INTO application_services_state(as_id, state, last_txn) "
+                "VALUES(?,?,?)"
+            ),
             (id, state, txn),
         )
 
     def _insert_txn(self, as_id, txn_id, events):
-        return self.db_pool.runQuery(
-            "INSERT INTO application_services_txns(as_id, txn_id, event_ids) "
-            "VALUES(?,?,?)",
+        return self.db_pool.runOperation(
+            self.engine.convert_param_style(
+                "INSERT INTO application_services_txns(as_id, txn_id, event_ids) "
+                "VALUES(?,?,?)"
+            ),
             (as_id, txn_id, json.dumps([e.event_id for e in events])),
         )
 
     def _set_last_txn(self, as_id, txn_id):
-        return self.db_pool.runQuery(
-            "INSERT INTO application_services_state(as_id, last_txn, state) "
-            "VALUES(?,?,?)",
+        return self.db_pool.runOperation(
+            self.engine.convert_param_style(
+                "INSERT INTO application_services_state(as_id, last_txn, state) "
+                "VALUES(?,?,?)"
+            ),
             (as_id, txn_id, ApplicationServiceState.UP),
         )
 
     @defer.inlineCallbacks
     def test_get_appservice_state_none(self):
-        service = Mock(id=999)
+        service = Mock(id="999")
         state = yield self.store.get_appservice_state(service)
         self.assertEquals(None, state)
 
@@ -200,7 +200,9 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         service = Mock(id=self.as_list[1]["id"])
         yield self.store.set_appservice_state(service, ApplicationServiceState.DOWN)
         rows = yield self.db_pool.runQuery(
-            "SELECT as_id FROM application_services_state WHERE state=?",
+            self.engine.convert_param_style(
+                "SELECT as_id FROM application_services_state WHERE state=?"
+            ),
             (ApplicationServiceState.DOWN,),
         )
         self.assertEquals(service.id, rows[0][0])
@@ -212,7 +214,9 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         yield self.store.set_appservice_state(service, ApplicationServiceState.DOWN)
         yield self.store.set_appservice_state(service, ApplicationServiceState.UP)
         rows = yield self.db_pool.runQuery(
-            "SELECT as_id FROM application_services_state WHERE state=?",
+            self.engine.convert_param_style(
+                "SELECT as_id FROM application_services_state WHERE state=?"
+            ),
             (ApplicationServiceState.UP,),
         )
         self.assertEquals(service.id, rows[0][0])
@@ -279,14 +283,19 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         yield self.store.complete_appservice_txn(txn_id=txn_id, service=service)
 
         res = yield self.db_pool.runQuery(
-            "SELECT last_txn FROM application_services_state WHERE as_id=?",
+            self.engine.convert_param_style(
+                "SELECT last_txn FROM application_services_state WHERE as_id=?"
+            ),
             (service.id,),
         )
         self.assertEquals(1, len(res))
         self.assertEquals(txn_id, res[0][0])
 
         res = yield self.db_pool.runQuery(
-            "SELECT * FROM application_services_txns WHERE txn_id=?", (txn_id,)
+            self.engine.convert_param_style(
+                "SELECT * FROM application_services_txns WHERE txn_id=?"
+            ),
+            (txn_id,),
         )
         self.assertEquals(0, len(res))
 
@@ -300,7 +309,9 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         yield self.store.complete_appservice_txn(txn_id=txn_id, service=service)
 
         res = yield self.db_pool.runQuery(
-            "SELECT last_txn, state FROM application_services_state WHERE " "as_id=?",
+            self.engine.convert_param_style(
+                "SELECT last_txn, state FROM application_services_state WHERE as_id=?"
+            ),
             (service.id,),
         )
         self.assertEquals(1, len(res))
@@ -308,7 +319,10 @@ class ApplicationServiceTransactionStoreTestCase(unittest.TestCase):
         self.assertEquals(ApplicationServiceState.UP, res[0][1])
 
         res = yield self.db_pool.runQuery(
-            "SELECT * FROM application_services_txns WHERE txn_id=?", (txn_id,)
+            self.engine.convert_param_style(
+                "SELECT * FROM application_services_txns WHERE txn_id=?"
+            ),
+            (txn_id,),
         )
         self.assertEquals(0, len(res))
 
@@ -394,37 +408,31 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
         f1 = self._write_config(suffix="1")
         f2 = self._write_config(suffix="2")
 
-        config = Mock(
-            app_service_config_files=[f1, f2], event_cache_size=1, password_providers=[]
-        )
         hs = yield setup_test_homeserver(
-            self.addCleanup,
-            config=config,
-            datastore=Mock(),
-            federation_sender=Mock(),
-            federation_client=Mock(),
+            self.addCleanup, federation_sender=Mock(), federation_client=Mock()
         )
 
-        ApplicationServiceStore(None, hs)
+        hs.config.app_service_config_files = [f1, f2]
+        hs.config.event_cache_size = 1
+        hs.config.password_providers = []
+
+        ApplicationServiceStore(hs.get_db_conn(), hs)
 
     @defer.inlineCallbacks
     def test_duplicate_ids(self):
         f1 = self._write_config(id="id", suffix="1")
         f2 = self._write_config(id="id", suffix="2")
 
-        config = Mock(
-            app_service_config_files=[f1, f2], event_cache_size=1, password_providers=[]
-        )
         hs = yield setup_test_homeserver(
-            self.addCleanup,
-            config=config,
-            datastore=Mock(),
-            federation_sender=Mock(),
-            federation_client=Mock(),
+            self.addCleanup, federation_sender=Mock(), federation_client=Mock()
         )
 
+        hs.config.app_service_config_files = [f1, f2]
+        hs.config.event_cache_size = 1
+        hs.config.password_providers = []
+
         with self.assertRaises(ConfigError) as cm:
-            ApplicationServiceStore(None, hs)
+            ApplicationServiceStore(hs.get_db_conn(), hs)
 
         e = cm.exception
         self.assertIn(f1, str(e))
@@ -436,19 +444,16 @@ class ApplicationServiceStoreConfigTestCase(unittest.TestCase):
         f1 = self._write_config(as_token="as_token", suffix="1")
         f2 = self._write_config(as_token="as_token", suffix="2")
 
-        config = Mock(
-            app_service_config_files=[f1, f2], event_cache_size=1, password_providers=[]
-        )
         hs = yield setup_test_homeserver(
-            self.addCleanup,
-            config=config,
-            datastore=Mock(),
-            federation_sender=Mock(),
-            federation_client=Mock(),
+            self.addCleanup, federation_sender=Mock(), federation_client=Mock()
         )
 
+        hs.config.app_service_config_files = [f1, f2]
+        hs.config.event_cache_size = 1
+        hs.config.password_providers = []
+
         with self.assertRaises(ConfigError) as cm:
-            ApplicationServiceStore(None, hs)
+            ApplicationServiceStore(hs.get_db_conn(), hs)
 
         e = cm.exception
         self.assertIn(f1, str(e))
diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py
index 7cb5f0e4cf..829f47d2e8 100644
--- a/tests/storage/test_base.py
+++ b/tests/storage/test_base.py
@@ -20,11 +20,11 @@ from mock import Mock
 
 from twisted.internet import defer
 
-from synapse.server import HomeServer
 from synapse.storage._base import SQLBaseStore
 from synapse.storage.engines import create_engine
 
 from tests import unittest
+from tests.utils import TestHomeServer
 
 
 class SQLBaseStoreTestCase(unittest.TestCase):
@@ -51,7 +51,7 @@ class SQLBaseStoreTestCase(unittest.TestCase):
         config = Mock()
         config.event_cache_size = 1
         config.database_config = {"name": "sqlite3"}
-        hs = HomeServer(
+        hs = TestHomeServer(
             "test",
             db_pool=self.db_pool,
             config=config,
diff --git a/tests/storage/test_client_ips.py b/tests/storage/test_client_ips.py
index c2e88bdbaf..4577e9422b 100644
--- a/tests/storage/test_client_ips.py
+++ b/tests/storage/test_client_ips.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 # Copyright 2016 OpenMarket Ltd
+# Copyright 2018 New Vector Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,35 +13,41 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 from mock import Mock
 
 from twisted.internet import defer
 
-import tests.unittest
-import tests.utils
+from synapse.http.site import XForwardedForRequest
+from synapse.rest.client.v1 import admin, login
+
+from tests import unittest
 
 
-class ClientIpStoreTestCase(tests.unittest.TestCase):
-    def __init__(self, *args, **kwargs):
-        super(ClientIpStoreTestCase, self).__init__(*args, **kwargs)
-        self.store = None  # type: synapse.storage.DataStore
-        self.clock = None  # type: tests.utils.MockClock
+class ClientIpStoreTestCase(unittest.HomeserverTestCase):
+    def make_homeserver(self, reactor, clock):
+        hs = self.setup_test_homeserver()
+        return hs
 
-    @defer.inlineCallbacks
-    def setUp(self):
-        self.hs = yield tests.utils.setup_test_homeserver(self.addCleanup)
+    def prepare(self, hs, reactor, clock):
         self.store = self.hs.get_datastore()
-        self.clock = self.hs.get_clock()
 
-    @defer.inlineCallbacks
     def test_insert_new_client_ip(self):
-        self.clock.now = 12345678
+        self.reactor.advance(12345678)
+
         user_id = "@user:id"
-        yield self.store.insert_client_ip(
-            user_id, "access_token", "ip", "user_agent", "device_id"
+        self.get_success(
+            self.store.insert_client_ip(
+                user_id, "access_token", "ip", "user_agent", "device_id"
+            )
         )
 
-        result = yield self.store.get_last_client_ip_by_device(user_id, "device_id")
+        # Trigger the storage loop
+        self.reactor.advance(10)
+
+        result = self.get_success(
+            self.store.get_last_client_ip_by_device(user_id, "device_id")
+        )
 
         r = result[(user_id, "device_id")]
         self.assertDictContainsSubset(
@@ -55,18 +62,18 @@ class ClientIpStoreTestCase(tests.unittest.TestCase):
             r,
         )
 
-    @defer.inlineCallbacks
     def test_disabled_monthly_active_user(self):
         self.hs.config.limit_usage_by_mau = False
         self.hs.config.max_mau_value = 50
         user_id = "@user:server"
-        yield self.store.insert_client_ip(
-            user_id, "access_token", "ip", "user_agent", "device_id"
+        self.get_success(
+            self.store.insert_client_ip(
+                user_id, "access_token", "ip", "user_agent", "device_id"
+            )
         )
-        active = yield self.store.user_last_seen_monthly_active(user_id)
+        active = self.get_success(self.store.user_last_seen_monthly_active(user_id))
         self.assertFalse(active)
 
-    @defer.inlineCallbacks
     def test_adding_monthly_active_user_when_full(self):
         self.hs.config.limit_usage_by_mau = True
         self.hs.config.max_mau_value = 50
@@ -76,40 +83,119 @@ class ClientIpStoreTestCase(tests.unittest.TestCase):
         self.store.get_monthly_active_count = Mock(
             return_value=defer.succeed(lots_of_users)
         )
-        yield self.store.insert_client_ip(
-            user_id, "access_token", "ip", "user_agent", "device_id"
+        self.get_success(
+            self.store.insert_client_ip(
+                user_id, "access_token", "ip", "user_agent", "device_id"
+            )
         )
-        active = yield self.store.user_last_seen_monthly_active(user_id)
+        active = self.get_success(self.store.user_last_seen_monthly_active(user_id))
         self.assertFalse(active)
 
-    @defer.inlineCallbacks
     def test_adding_monthly_active_user_when_space(self):
         self.hs.config.limit_usage_by_mau = True
         self.hs.config.max_mau_value = 50
         user_id = "@user:server"
-        active = yield self.store.user_last_seen_monthly_active(user_id)
+        active = self.get_success(self.store.user_last_seen_monthly_active(user_id))
         self.assertFalse(active)
 
-        yield self.store.insert_client_ip(
-            user_id, "access_token", "ip", "user_agent", "device_id"
+        # Trigger the saving loop
+        self.reactor.advance(10)
+
+        self.get_success(
+            self.store.insert_client_ip(
+                user_id, "access_token", "ip", "user_agent", "device_id"
+            )
         )
-        active = yield self.store.user_last_seen_monthly_active(user_id)
+        active = self.get_success(self.store.user_last_seen_monthly_active(user_id))
         self.assertTrue(active)
 
-    @defer.inlineCallbacks
     def test_updating_monthly_active_user_when_space(self):
         self.hs.config.limit_usage_by_mau = True
         self.hs.config.max_mau_value = 50
         user_id = "@user:server"
+        self.get_success(
+            self.store.register(user_id=user_id, token="123", password_hash=None)
+        )
 
-        active = yield self.store.user_last_seen_monthly_active(user_id)
+        active = self.get_success(self.store.user_last_seen_monthly_active(user_id))
         self.assertFalse(active)
 
-        yield self.store.insert_client_ip(
-            user_id, "access_token", "ip", "user_agent", "device_id"
-        )
-        yield self.store.insert_client_ip(
-            user_id, "access_token", "ip", "user_agent", "device_id"
+        # Trigger the saving loop
+        self.reactor.advance(10)
+
+        self.get_success(
+            self.store.insert_client_ip(
+                user_id, "access_token", "ip", "user_agent", "device_id"
+            )
         )
-        active = yield self.store.user_last_seen_monthly_active(user_id)
+        active = self.get_success(self.store.user_last_seen_monthly_active(user_id))
         self.assertTrue(active)
+
+
+class ClientIpAuthTestCase(unittest.HomeserverTestCase):
+
+    servlets = [admin.register_servlets, login.register_servlets]
+
+    def make_homeserver(self, reactor, clock):
+        hs = self.setup_test_homeserver()
+        return hs
+
+    def prepare(self, hs, reactor, clock):
+        self.store = self.hs.get_datastore()
+        self.user_id = self.register_user("bob", "abc123", True)
+
+    def test_request_with_xforwarded(self):
+        """
+        The IP in X-Forwarded-For is entered into the client IPs table.
+        """
+        self._runtest(
+            {b"X-Forwarded-For": b"127.9.0.1"},
+            "127.9.0.1",
+            {"request": XForwardedForRequest},
+        )
+
+    def test_request_from_getPeer(self):
+        """
+        The IP returned by getPeer is entered into the client IPs table, if
+        there's no X-Forwarded-For header.
+        """
+        self._runtest({}, "127.0.0.1", {})
+
+    def _runtest(self, headers, expected_ip, make_request_args):
+        device_id = "bleb"
+
+        access_token = self.login("bob", "abc123", device_id=device_id)
+
+        # Advance to a known time
+        self.reactor.advance(123456 - self.reactor.seconds())
+
+        request, channel = self.make_request(
+            "GET",
+            "/_matrix/client/r0/admin/users/" + self.user_id,
+            access_token=access_token,
+            **make_request_args
+        )
+        request.requestHeaders.addRawHeader(b"User-Agent", b"Mozzila pizza")
+
+        # Add the optional headers
+        for h, v in headers.items():
+            request.requestHeaders.addRawHeader(h, v)
+        self.render(request)
+
+        # Advance so the save loop occurs
+        self.reactor.advance(100)
+
+        result = self.get_success(
+            self.store.get_last_client_ip_by_device(self.user_id, device_id)
+        )
+        r = result[(self.user_id, device_id)]
+        self.assertDictContainsSubset(
+            {
+                "user_id": self.user_id,
+                "device_id": device_id,
+                "ip": expected_ip,
+                "user_agent": "Mozzila pizza",
+                "last_seen": 123456100,
+            },
+            r,
+        )
diff --git a/tests/storage/test_directory.py b/tests/storage/test_directory.py
index b4510c1c8d..4e128e1047 100644
--- a/tests/storage/test_directory.py
+++ b/tests/storage/test_directory.py
@@ -16,7 +16,6 @@
 
 from twisted.internet import defer
 
-from synapse.storage.directory import DirectoryStore
 from synapse.types import RoomAlias, RoomID
 
 from tests import unittest
@@ -28,7 +27,7 @@ class DirectoryStoreTestCase(unittest.TestCase):
     def setUp(self):
         hs = yield setup_test_homeserver(self.addCleanup)
 
-        self.store = DirectoryStore(None, hs)
+        self.store = hs.get_datastore()
 
         self.room = RoomID.from_string("!abcde:test")
         self.alias = RoomAlias.from_string("#my-room:test")
diff --git a/tests/storage/test_event_federation.py b/tests/storage/test_event_federation.py
index 2fdf34fdf6..0d4e74d637 100644
--- a/tests/storage/test_event_federation.py
+++ b/tests/storage/test_event_federation.py
@@ -37,10 +37,10 @@ class EventFederationWorkerStoreTestCase(tests.unittest.TestCase):
                 (
                     "INSERT INTO events ("
                     "   room_id, event_id, type, depth, topological_ordering,"
-                    "   content, processed, outlier) "
-                    "VALUES (?, ?, 'm.test', ?, ?, 'test', ?, ?)"
+                    "   content, processed, outlier, stream_ordering) "
+                    "VALUES (?, ?, 'm.test', ?, ?, 'test', ?, ?, ?)"
                 ),
-                (room_id, event_id, i, i, True, False),
+                (room_id, event_id, i, i, True, False, i),
             )
 
             txn.execute(
diff --git a/tests/storage/test_monthly_active_users.py b/tests/storage/test_monthly_active_users.py
index f2ed866ae7..686f12a0dc 100644
--- a/tests/storage/test_monthly_active_users.py
+++ b/tests/storage/test_monthly_active_users.py
@@ -12,26 +12,27 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from mock import Mock
 
 from twisted.internet import defer
 
-import tests.unittest
-import tests.utils
-from tests.utils import setup_test_homeserver
+from tests.unittest import HomeserverTestCase
 
 FORTY_DAYS = 40 * 24 * 60 * 60
 
 
-class MonthlyActiveUsersTestCase(tests.unittest.TestCase):
-    def __init__(self, *args, **kwargs):
-        super(MonthlyActiveUsersTestCase, self).__init__(*args, **kwargs)
+class MonthlyActiveUsersTestCase(HomeserverTestCase):
+    def make_homeserver(self, reactor, clock):
 
-    @defer.inlineCallbacks
-    def setUp(self):
-        self.hs = yield setup_test_homeserver(self.addCleanup)
-        self.store = self.hs.get_datastore()
+        hs = self.setup_test_homeserver()
+        self.store = hs.get_datastore()
+        hs.config.limit_usage_by_mau = True
+        hs.config.max_mau_value = 50
+        # Advance the clock a bit
+        reactor.advance(FORTY_DAYS)
+
+        return hs
 
-    @defer.inlineCallbacks
     def test_initialise_reserved_users(self):
         self.hs.config.max_mau_value = 5
         user1 = "@user1:server"
@@ -44,88 +45,172 @@ class MonthlyActiveUsersTestCase(tests.unittest.TestCase):
         ]
         user_num = len(threepids)
 
-        yield self.store.register(user_id=user1, token="123", password_hash=None)
-
-        yield self.store.register(user_id=user2, token="456", password_hash=None)
+        self.store.register(user_id=user1, token="123", password_hash=None)
+        self.store.register(user_id=user2, token="456", password_hash=None)
+        self.pump()
 
         now = int(self.hs.get_clock().time_msec())
-        yield self.store.user_add_threepid(user1, "email", user1_email, now, now)
-        yield self.store.user_add_threepid(user2, "email", user2_email, now, now)
-        yield self.store.initialise_reserved_users(threepids)
+        self.store.user_add_threepid(user1, "email", user1_email, now, now)
+        self.store.user_add_threepid(user2, "email", user2_email, now, now)
+        self.store.initialise_reserved_users(threepids)
+        self.pump()
 
-        active_count = yield self.store.get_monthly_active_count()
+        active_count = self.store.get_monthly_active_count()
 
         # Test total counts
-        self.assertEquals(active_count, user_num)
+        self.assertEquals(self.get_success(active_count), user_num)
 
         # Test user is marked as active
-
-        timestamp = yield self.store.user_last_seen_monthly_active(user1)
-        self.assertTrue(timestamp)
-        timestamp = yield self.store.user_last_seen_monthly_active(user2)
-        self.assertTrue(timestamp)
+        timestamp = self.store.user_last_seen_monthly_active(user1)
+        self.assertTrue(self.get_success(timestamp))
+        timestamp = self.store.user_last_seen_monthly_active(user2)
+        self.assertTrue(self.get_success(timestamp))
 
         # Test that users are never removed from the db.
         self.hs.config.max_mau_value = 0
 
-        self.hs.get_clock().advance_time(FORTY_DAYS)
+        self.reactor.advance(FORTY_DAYS)
 
-        yield self.store.reap_monthly_active_users()
+        self.store.reap_monthly_active_users()
+        self.pump()
 
-        active_count = yield self.store.get_monthly_active_count()
-        self.assertEquals(active_count, user_num)
+        active_count = self.store.get_monthly_active_count()
+        self.assertEquals(self.get_success(active_count), user_num)
 
-        # Test that regalar users are removed from the db
+        # Test that regular users are removed from the db
         ru_count = 2
-        yield self.store.upsert_monthly_active_user("@ru1:server")
-        yield self.store.upsert_monthly_active_user("@ru2:server")
-        active_count = yield self.store.get_monthly_active_count()
+        self.store.upsert_monthly_active_user("@ru1:server")
+        self.store.upsert_monthly_active_user("@ru2:server")
+        self.pump()
 
-        self.assertEqual(active_count, user_num + ru_count)
+        active_count = self.store.get_monthly_active_count()
+        self.assertEqual(self.get_success(active_count), user_num + ru_count)
         self.hs.config.max_mau_value = user_num
-        yield self.store.reap_monthly_active_users()
+        self.store.reap_monthly_active_users()
+        self.pump()
 
-        active_count = yield self.store.get_monthly_active_count()
-        self.assertEquals(active_count, user_num)
+        active_count = self.store.get_monthly_active_count()
+        self.assertEquals(self.get_success(active_count), user_num)
 
-    @defer.inlineCallbacks
     def test_can_insert_and_count_mau(self):
-        count = yield self.store.get_monthly_active_count()
-        self.assertEqual(0, count)
+        count = self.store.get_monthly_active_count()
+        self.assertEqual(0, self.get_success(count))
 
-        yield self.store.upsert_monthly_active_user("@user:server")
-        count = yield self.store.get_monthly_active_count()
+        self.store.upsert_monthly_active_user("@user:server")
+        self.pump()
 
-        self.assertEqual(1, count)
+        count = self.store.get_monthly_active_count()
+        self.assertEqual(1, self.get_success(count))
 
-    @defer.inlineCallbacks
     def test_user_last_seen_monthly_active(self):
         user_id1 = "@user1:server"
         user_id2 = "@user2:server"
         user_id3 = "@user3:server"
 
-        result = yield self.store.user_last_seen_monthly_active(user_id1)
-        self.assertFalse(result == 0)
-        yield self.store.upsert_monthly_active_user(user_id1)
-        yield self.store.upsert_monthly_active_user(user_id2)
-        result = yield self.store.user_last_seen_monthly_active(user_id1)
-        self.assertTrue(result > 0)
-        result = yield self.store.user_last_seen_monthly_active(user_id3)
-        self.assertFalse(result == 0)
+        result = self.store.user_last_seen_monthly_active(user_id1)
+        self.assertFalse(self.get_success(result) == 0)
+
+        self.store.upsert_monthly_active_user(user_id1)
+        self.store.upsert_monthly_active_user(user_id2)
+        self.pump()
+
+        result = self.store.user_last_seen_monthly_active(user_id1)
+        self.assertGreater(self.get_success(result), 0)
+
+        result = self.store.user_last_seen_monthly_active(user_id3)
+        self.assertNotEqual(self.get_success(result), 0)
 
-    @defer.inlineCallbacks
     def test_reap_monthly_active_users(self):
         self.hs.config.max_mau_value = 5
         initial_users = 10
         for i in range(initial_users):
-            yield self.store.upsert_monthly_active_user("@user%d:server" % i)
-        count = yield self.store.get_monthly_active_count()
-        self.assertTrue(count, initial_users)
-        yield self.store.reap_monthly_active_users()
-        count = yield self.store.get_monthly_active_count()
-        self.assertEquals(count, initial_users - self.hs.config.max_mau_value)
-
-        self.hs.get_clock().advance_time(FORTY_DAYS)
-        yield self.store.reap_monthly_active_users()
-        count = yield self.store.get_monthly_active_count()
-        self.assertEquals(count, 0)
+            self.store.upsert_monthly_active_user("@user%d:server" % i)
+        self.pump()
+
+        count = self.store.get_monthly_active_count()
+        self.assertTrue(self.get_success(count), initial_users)
+
+        self.store.reap_monthly_active_users()
+        self.pump()
+        count = self.store.get_monthly_active_count()
+        self.assertEquals(
+            self.get_success(count), initial_users - self.hs.config.max_mau_value
+        )
+
+        self.reactor.advance(FORTY_DAYS)
+        self.store.reap_monthly_active_users()
+        self.pump()
+
+        count = self.store.get_monthly_active_count()
+        self.assertEquals(self.get_success(count), 0)
+
+    def test_populate_monthly_users_is_guest(self):
+        # Test that guest users are not added to mau list
+        user_id = "user_id"
+        self.store.register(
+            user_id=user_id, token="123", password_hash=None, make_guest=True
+        )
+        self.store.upsert_monthly_active_user = Mock()
+        self.store.populate_monthly_active_users(user_id)
+        self.pump()
+        self.store.upsert_monthly_active_user.assert_not_called()
+
+    def test_populate_monthly_users_should_update(self):
+        self.store.upsert_monthly_active_user = Mock()
+
+        self.store.is_trial_user = Mock(
+            return_value=defer.succeed(False)
+        )
+
+        self.store.user_last_seen_monthly_active = Mock(
+            return_value=defer.succeed(None)
+        )
+        self.store.populate_monthly_active_users('user_id')
+        self.pump()
+        self.store.upsert_monthly_active_user.assert_called_once()
+
+    def test_populate_monthly_users_should_not_update(self):
+        self.store.upsert_monthly_active_user = Mock()
+
+        self.store.is_trial_user = Mock(
+            return_value=defer.succeed(False)
+        )
+        self.store.user_last_seen_monthly_active = Mock(
+            return_value=defer.succeed(
+                self.hs.get_clock().time_msec()
+            )
+        )
+        self.store.populate_monthly_active_users('user_id')
+        self.pump()
+        self.store.upsert_monthly_active_user.assert_not_called()
+
+    def test_get_reserved_real_user_account(self):
+        # Test no reserved users, or reserved threepids
+        count = self.store.get_registered_reserved_users_count()
+        self.assertEquals(self.get_success(count), 0)
+        # Test reserved users but no registered users
+
+        user1 = '@user1:example.com'
+        user2 = '@user2:example.com'
+        user1_email = 'user1@example.com'
+        user2_email = 'user2@example.com'
+        threepids = [
+            {'medium': 'email', 'address': user1_email},
+            {'medium': 'email', 'address': user2_email},
+        ]
+        self.hs.config.mau_limits_reserved_threepids = threepids
+        self.store.initialise_reserved_users(threepids)
+        self.pump()
+        count = self.store.get_registered_reserved_users_count()
+        self.assertEquals(self.get_success(count), 0)
+
+        # Test reserved registed users
+        self.store.register(user_id=user1, token="123", password_hash=None)
+        self.store.register(user_id=user2, token="456", password_hash=None)
+        self.pump()
+
+        now = int(self.hs.get_clock().time_msec())
+        self.store.user_add_threepid(user1, "email", user1_email, now, now)
+        self.store.user_add_threepid(user2, "email", user2_email, now, now)
+        count = self.store.get_registered_reserved_users_count()
+        self.assertEquals(self.get_success(count), len(threepids))
diff --git a/tests/storage/test_presence.py b/tests/storage/test_presence.py
index b5b58ff660..c7a63f39b9 100644
--- a/tests/storage/test_presence.py
+++ b/tests/storage/test_presence.py
@@ -16,19 +16,18 @@
 
 from twisted.internet import defer
 
-from synapse.storage.presence import PresenceStore
 from synapse.types import UserID
 
 from tests import unittest
-from tests.utils import MockClock, setup_test_homeserver
+from tests.utils import setup_test_homeserver
 
 
 class PresenceStoreTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def setUp(self):
-        hs = yield setup_test_homeserver(self.addCleanup, clock=MockClock())
+        hs = yield setup_test_homeserver(self.addCleanup)
 
-        self.store = PresenceStore(None, hs)
+        self.store = hs.get_datastore()
 
         self.u_apple = UserID.from_string("@apple:test")
         self.u_banana = UserID.from_string("@banana:test")
diff --git a/tests/storage/test_profile.py b/tests/storage/test_profile.py
index a1f6618bf9..45824bd3b2 100644
--- a/tests/storage/test_profile.py
+++ b/tests/storage/test_profile.py
@@ -28,7 +28,7 @@ class ProfileStoreTestCase(unittest.TestCase):
     def setUp(self):
         hs = yield setup_test_homeserver(self.addCleanup)
 
-        self.store = ProfileStore(None, hs)
+        self.store = ProfileStore(hs.get_db_conn(), hs)
 
         self.u_frank = UserID.from_string("@frank:test")
 
diff --git a/tests/storage/test_purge.py b/tests/storage/test_purge.py
new file mode 100644
index 0000000000..f671599cb8
--- /dev/null
+++ b/tests/storage/test_purge.py
@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 New Vector Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from synapse.rest.client.v1 import room
+
+from tests.unittest import HomeserverTestCase
+
+
+class PurgeTests(HomeserverTestCase):
+
+    user_id = "@red:server"
+    servlets = [room.register_servlets]
+
+    def make_homeserver(self, reactor, clock):
+        hs = self.setup_test_homeserver("server", http_client=None)
+        return hs
+
+    def prepare(self, reactor, clock, hs):
+        self.room_id = self.helper.create_room_as(self.user_id)
+
+    def test_purge(self):
+        """
+        Purging a room will delete everything before the topological point.
+        """
+        # Send four messages to the room
+        first = self.helper.send(self.room_id, body="test1")
+        second = self.helper.send(self.room_id, body="test2")
+        third = self.helper.send(self.room_id, body="test3")
+        last = self.helper.send(self.room_id, body="test4")
+
+        storage = self.hs.get_datastore()
+
+        # Get the topological token
+        event = storage.get_topological_token_for_event(last["event_id"])
+        self.pump()
+        event = self.successResultOf(event)
+
+        # Purge everything before this topological token
+        purge = storage.purge_history(self.room_id, event, True)
+        self.pump()
+        self.assertEqual(self.successResultOf(purge), None)
+
+        # Try and get the events
+        get_first = storage.get_event(first["event_id"])
+        get_second = storage.get_event(second["event_id"])
+        get_third = storage.get_event(third["event_id"])
+        get_last = storage.get_event(last["event_id"])
+        self.pump()
+
+        # 1-3 should fail and last will succeed, meaning that 1-3 are deleted
+        # and last is not.
+        self.failureResultOf(get_first)
+        self.failureResultOf(get_second)
+        self.failureResultOf(get_third)
+        self.successResultOf(get_last)
+
+    def test_purge_wont_delete_extrems(self):
+        """
+        Purging a room will delete everything before the topological point.
+        """
+        # Send four messages to the room
+        first = self.helper.send(self.room_id, body="test1")
+        second = self.helper.send(self.room_id, body="test2")
+        third = self.helper.send(self.room_id, body="test3")
+        last = self.helper.send(self.room_id, body="test4")
+
+        storage = self.hs.get_datastore()
+
+        # Set the topological token higher than it should be
+        event = storage.get_topological_token_for_event(last["event_id"])
+        self.pump()
+        event = self.successResultOf(event)
+        event = "t{}-{}".format(
+            *list(map(lambda x: x + 1, map(int, event[1:].split("-"))))
+        )
+
+        # Purge everything before this topological token
+        purge = storage.purge_history(self.room_id, event, True)
+        self.pump()
+        f = self.failureResultOf(purge)
+        self.assertIn("greater than forward", f.value.args[0])
+
+        # Try and get the events
+        get_first = storage.get_event(first["event_id"])
+        get_second = storage.get_event(second["event_id"])
+        get_third = storage.get_event(third["event_id"])
+        get_last = storage.get_event(last["event_id"])
+        self.pump()
+
+        # Nothing is deleted.
+        self.successResultOf(get_first)
+        self.successResultOf(get_second)
+        self.successResultOf(get_third)
+        self.successResultOf(get_last)
diff --git a/tests/storage/test_state.py b/tests/storage/test_state.py
index d717b9f94e..b9c5b39d59 100644
--- a/tests/storage/test_state.py
+++ b/tests/storage/test_state.py
@@ -75,6 +75,45 @@ class StateStoreTestCase(tests.unittest.TestCase):
         self.assertEqual(len(s1), len(s2))
 
     @defer.inlineCallbacks
+    def test_get_state_groups_ids(self):
+        e1 = yield self.inject_state_event(
+            self.room, self.u_alice, EventTypes.Create, '', {}
+        )
+        e2 = yield self.inject_state_event(
+            self.room, self.u_alice, EventTypes.Name, '', {"name": "test room"}
+        )
+
+        state_group_map = yield self.store.get_state_groups_ids(self.room, [e2.event_id])
+        self.assertEqual(len(state_group_map), 1)
+        state_map = list(state_group_map.values())[0]
+        self.assertDictEqual(
+            state_map,
+            {
+                (EventTypes.Create, ''): e1.event_id,
+                (EventTypes.Name, ''): e2.event_id,
+            },
+        )
+
+    @defer.inlineCallbacks
+    def test_get_state_groups(self):
+        e1 = yield self.inject_state_event(
+            self.room, self.u_alice, EventTypes.Create, '', {}
+        )
+        e2 = yield self.inject_state_event(
+            self.room, self.u_alice, EventTypes.Name, '', {"name": "test room"}
+        )
+
+        state_group_map = yield self.store.get_state_groups(
+            self.room, [e2.event_id])
+        self.assertEqual(len(state_group_map), 1)
+        state_list = list(state_group_map.values())[0]
+
+        self.assertEqual(
+            {ev.event_id for ev in state_list},
+            {e1.event_id, e2.event_id},
+        )
+
+    @defer.inlineCallbacks
     def test_get_state_for_event(self):
 
         # this defaults to a linear DAG as each new injection defaults to whatever
@@ -185,8 +224,7 @@ class StateStoreTestCase(tests.unittest.TestCase):
 
         # test _get_some_state_from_cache correctly filters out members with types=[]
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
-            self.store._state_group_cache,
-            group, [], filtered_types=[EventTypes.Member]
+            self.store._state_group_cache, group, [], filtered_types=[EventTypes.Member]
         )
 
         self.assertEqual(is_all, True)
@@ -200,19 +238,20 @@ class StateStoreTestCase(tests.unittest.TestCase):
 
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_members_cache,
-            group, [], filtered_types=[EventTypes.Member]
+            group,
+            [],
+            filtered_types=[EventTypes.Member],
         )
 
         self.assertEqual(is_all, True)
-        self.assertDictEqual(
-            {},
-            state_dict,
-        )
+        self.assertDictEqual({}, state_dict)
 
         # test _get_some_state_from_cache correctly filters in members with wildcard types
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_cache,
-            group, [(EventTypes.Member, None)], filtered_types=[EventTypes.Member]
+            group,
+            [(EventTypes.Member, None)],
+            filtered_types=[EventTypes.Member],
         )
 
         self.assertEqual(is_all, True)
@@ -226,7 +265,9 @@ class StateStoreTestCase(tests.unittest.TestCase):
 
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_members_cache,
-            group, [(EventTypes.Member, None)], filtered_types=[EventTypes.Member]
+            group,
+            [(EventTypes.Member, None)],
+            filtered_types=[EventTypes.Member],
         )
 
         self.assertEqual(is_all, True)
@@ -264,18 +305,15 @@ class StateStoreTestCase(tests.unittest.TestCase):
         )
 
         self.assertEqual(is_all, True)
-        self.assertDictEqual(
-            {
-                (e5.type, e5.state_key): e5.event_id,
-            },
-            state_dict,
-        )
+        self.assertDictEqual({(e5.type, e5.state_key): e5.event_id}, state_dict)
 
         # test _get_some_state_from_cache correctly filters in members with specific types
         # and no filtered_types
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_members_cache,
-            group, [(EventTypes.Member, e5.state_key)], filtered_types=None
+            group,
+            [(EventTypes.Member, e5.state_key)],
+            filtered_types=None,
         )
 
         self.assertEqual(is_all, True)
@@ -305,9 +343,7 @@ class StateStoreTestCase(tests.unittest.TestCase):
             key=group,
             value=state_dict_ids,
             # list fetched keys so it knows it's partial
-            fetched_keys=(
-                (e1.type, e1.state_key),
-            ),
+            fetched_keys=((e1.type, e1.state_key),),
         )
 
         (is_all, known_absent, state_dict_ids) = self.store._state_group_cache.get(
@@ -315,20 +351,8 @@ class StateStoreTestCase(tests.unittest.TestCase):
         )
 
         self.assertEqual(is_all, False)
-        self.assertEqual(
-            known_absent,
-            set(
-                [
-                    (e1.type, e1.state_key),
-                ]
-            ),
-        )
-        self.assertDictEqual(
-            state_dict_ids,
-            {
-                (e1.type, e1.state_key): e1.event_id,
-            },
-        )
+        self.assertEqual(known_absent, set([(e1.type, e1.state_key)]))
+        self.assertDictEqual(state_dict_ids, {(e1.type, e1.state_key): e1.event_id})
 
         ############################################
         # test that things work with a partial cache
@@ -336,8 +360,7 @@ class StateStoreTestCase(tests.unittest.TestCase):
         # test _get_some_state_from_cache correctly filters out members with types=[]
         room_id = self.room.to_string()
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
-            self.store._state_group_cache,
-            group, [], filtered_types=[EventTypes.Member]
+            self.store._state_group_cache, group, [], filtered_types=[EventTypes.Member]
         )
 
         self.assertEqual(is_all, False)
@@ -346,7 +369,9 @@ class StateStoreTestCase(tests.unittest.TestCase):
         room_id = self.room.to_string()
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_members_cache,
-            group, [], filtered_types=[EventTypes.Member]
+            group,
+            [],
+            filtered_types=[EventTypes.Member],
         )
 
         self.assertEqual(is_all, True)
@@ -355,20 +380,19 @@ class StateStoreTestCase(tests.unittest.TestCase):
         # test _get_some_state_from_cache correctly filters in members wildcard types
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_cache,
-            group, [(EventTypes.Member, None)], filtered_types=[EventTypes.Member]
+            group,
+            [(EventTypes.Member, None)],
+            filtered_types=[EventTypes.Member],
         )
 
         self.assertEqual(is_all, False)
-        self.assertDictEqual(
-            {
-                (e1.type, e1.state_key): e1.event_id,
-            },
-            state_dict,
-        )
+        self.assertDictEqual({(e1.type, e1.state_key): e1.event_id}, state_dict)
 
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_members_cache,
-            group, [(EventTypes.Member, None)], filtered_types=[EventTypes.Member]
+            group,
+            [(EventTypes.Member, None)],
+            filtered_types=[EventTypes.Member],
         )
 
         self.assertEqual(is_all, True)
@@ -389,12 +413,7 @@ class StateStoreTestCase(tests.unittest.TestCase):
         )
 
         self.assertEqual(is_all, False)
-        self.assertDictEqual(
-            {
-                (e1.type, e1.state_key): e1.event_id,
-            },
-            state_dict,
-        )
+        self.assertDictEqual({(e1.type, e1.state_key): e1.event_id}, state_dict)
 
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_members_cache,
@@ -404,18 +423,15 @@ class StateStoreTestCase(tests.unittest.TestCase):
         )
 
         self.assertEqual(is_all, True)
-        self.assertDictEqual(
-            {
-                (e5.type, e5.state_key): e5.event_id,
-            },
-            state_dict,
-        )
+        self.assertDictEqual({(e5.type, e5.state_key): e5.event_id}, state_dict)
 
         # test _get_some_state_from_cache correctly filters in members with specific types
         # and no filtered_types
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_cache,
-            group, [(EventTypes.Member, e5.state_key)], filtered_types=None
+            group,
+            [(EventTypes.Member, e5.state_key)],
+            filtered_types=None,
         )
 
         self.assertEqual(is_all, False)
@@ -423,13 +439,10 @@ class StateStoreTestCase(tests.unittest.TestCase):
 
         (state_dict, is_all) = yield self.store._get_some_state_from_cache(
             self.store._state_group_members_cache,
-            group, [(EventTypes.Member, e5.state_key)], filtered_types=None
+            group,
+            [(EventTypes.Member, e5.state_key)],
+            filtered_types=None,
         )
 
         self.assertEqual(is_all, True)
-        self.assertDictEqual(
-            {
-                (e5.type, e5.state_key): e5.event_id,
-            },
-            state_dict,
-        )
+        self.assertDictEqual({(e5.type, e5.state_key): e5.event_id}, state_dict)
diff --git a/tests/storage/test_transactions.py b/tests/storage/test_transactions.py
new file mode 100644
index 0000000000..14169afa96
--- /dev/null
+++ b/tests/storage/test_transactions.py
@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 New Vector Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from tests.unittest import HomeserverTestCase
+
+
+class TransactionStoreTestCase(HomeserverTestCase):
+    def prepare(self, reactor, clock, homeserver):
+        self.store = homeserver.get_datastore()
+
+    def test_get_set_transactions(self):
+        """Tests that we can successfully get a non-existent entry for
+        destination retries, as well as testing tht we can set and get
+        correctly.
+        """
+        d = self.store.get_destination_retry_timings("example.com")
+        r = self.get_success(d)
+        self.assertIsNone(r)
+
+        d = self.store.set_destination_retry_timings("example.com", 50, 100)
+        self.get_success(d)
+
+        d = self.store.get_destination_retry_timings("example.com")
+        r = self.get_success(d)
+
+        self.assert_dict({"retry_last_ts": 50, "retry_interval": 100}, r)
+
+    def test_initial_set_transactions(self):
+        """Tests that we can successfully set the destination retries (there
+        was a bug around invalidating the cache that broke this)
+        """
+        d = self.store.set_destination_retry_timings("example.com", 50, 100)
+        self.get_success(d)
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index b46e0ea7e2..0dde1ab2fe 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -30,7 +30,7 @@ class UserDirectoryStoreTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def setUp(self):
         self.hs = yield setup_test_homeserver(self.addCleanup)
-        self.store = UserDirectoryStore(None, self.hs)
+        self.store = UserDirectoryStore(self.hs.get_db_conn(), self.hs)
 
         # alice and bob are both in !room_id. bobby is not but shares
         # a homeserver with alice.