From 414fa36f3e22d87ea6ce4bdd6af54bce5a0de67c Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Tue, 28 Aug 2018 17:21:05 +0100
Subject: [PATCH] Fix up tests

---
 tests/storage/test_base.py | 4 ++--
 tests/test_types.py        | 4 ++--
 tests/utils.py             | 8 ++++++--
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py
index 7cb5f0e4c..829f47d2e 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/test_types.py b/tests/test_types.py
index be072d402..0f5c8bfaf 100644
--- a/tests/test_types.py
+++ b/tests/test_types.py
@@ -14,12 +14,12 @@
 # limitations under the License.
 
 from synapse.api.errors import SynapseError
-from synapse.server import HomeServer
 from synapse.types import GroupID, RoomAlias, UserID
 
 from tests import unittest
+from tests.utils import TestHomeServer
 
-mock_homeserver = HomeServer(hostname="my.domain")
+mock_homeserver = TestHomeServer(hostname="my.domain")
 
 
 class UserIDTestCase(unittest.TestCase):
diff --git a/tests/utils.py b/tests/utils.py
index e8ef10445..657c258a6 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -29,7 +29,7 @@ from synapse.api.errors import CodeMessageException, cs_error
 from synapse.federation.transport import server
 from synapse.http.server import HttpServer
 from synapse.server import HomeServer
-from synapse.storage import PostgresEngine
+from synapse.storage import DataStore, PostgresEngine
 from synapse.storage.engines import create_engine
 from synapse.storage.prepare_database import (
     _get_or_create_schema_state,
@@ -92,10 +92,14 @@ def setupdb():
         atexit.register(_cleanup)
 
 
+class TestHomeServer(HomeServer):
+    DATASTORE_CLASS = DataStore
+
+
 @defer.inlineCallbacks
 def setup_test_homeserver(
     cleanup_func, name="test", datastore=None, config=None, reactor=None,
-    homeserverToUse=HomeServer, **kargs
+    homeserverToUse=TestHomeServer, **kargs
 ):
     """
     Setup a homeserver suitable for running tests against.  Keyword arguments