Merge branch 'develop' of github.com:matrix-org/synapse into erikj/unfederatable

11 files changed, 391 insertions, 98 deletions

diff --git a/tests/api/test_auth.py b/tests/api/test_auth.py
index 22fc804331..c96273480d 100644
--- a/tests/api/test_auth.py
+++ b/tests/api/test_auth.py
@@ -19,17 +19,21 @@ from mock import Mock
 
 from synapse.api.auth import Auth
 from synapse.api.errors import AuthError
+from synapse.types import UserID
+from tests.utils import setup_test_homeserver
+
+import pymacaroons
 
 
 class AuthTestCase(unittest.TestCase):
 
+    @defer.inlineCallbacks
     def setUp(self):
         self.state_handler = Mock()
         self.store = Mock()
 
-        self.hs = Mock()
+        self.hs = yield setup_test_homeserver(handlers=None)
         self.hs.get_datastore = Mock(return_value=self.store)
-        self.hs.get_state_handler = Mock(return_value=self.state_handler)
         self.auth = Auth(self.hs)
 
         self.test_user = "@foo:bar"
@@ -133,3 +137,140 @@ class AuthTestCase(unittest.TestCase):
         request.requestHeaders.getRawHeaders = Mock(return_value=[""])
         d = self.auth.get_user_by_req(request)
         self.failureResultOf(d, AuthError)
+
+    @defer.inlineCallbacks
+    def test_get_user_from_macaroon(self):
+        # TODO(danielwh): Remove this mock when we remove the
+        # get_user_by_access_token fallback.
+        self.store.get_user_by_access_token = Mock(
+            return_value={"name": "@baldrick:matrix.org"}
+        )
+
+        user_id = "@baldrick:matrix.org"
+        macaroon = pymacaroons.Macaroon(
+            location=self.hs.config.server_name,
+            identifier="key",
+            key=self.hs.config.macaroon_secret_key)
+        macaroon.add_first_party_caveat("gen = 1")
+        macaroon.add_first_party_caveat("type = access")
+        macaroon.add_first_party_caveat("user_id = %s" % (user_id,))
+        user_info = yield self.auth._get_user_from_macaroon(macaroon.serialize())
+        user = user_info["user"]
+        self.assertEqual(UserID.from_string(user_id), user)
+
+    @defer.inlineCallbacks
+    def test_get_user_from_macaroon_user_db_mismatch(self):
+        self.store.get_user_by_access_token = Mock(
+            return_value={"name": "@percy:matrix.org"}
+        )
+
+        user = "@baldrick:matrix.org"
+        macaroon = pymacaroons.Macaroon(
+            location=self.hs.config.server_name,
+            identifier="key",
+            key=self.hs.config.macaroon_secret_key)
+        macaroon.add_first_party_caveat("gen = 1")
+        macaroon.add_first_party_caveat("type = access")
+        macaroon.add_first_party_caveat("user_id = %s" % (user,))
+        with self.assertRaises(AuthError) as cm:
+            yield self.auth._get_user_from_macaroon(macaroon.serialize())
+        self.assertEqual(401, cm.exception.code)
+        self.assertIn("User mismatch", cm.exception.msg)
+
+    @defer.inlineCallbacks
+    def test_get_user_from_macaroon_missing_caveat(self):
+        # TODO(danielwh): Remove this mock when we remove the
+        # get_user_by_access_token fallback.
+        self.store.get_user_by_access_token = Mock(
+            return_value={"name": "@baldrick:matrix.org"}
+        )
+
+        macaroon = pymacaroons.Macaroon(
+            location=self.hs.config.server_name,
+            identifier="key",
+            key=self.hs.config.macaroon_secret_key)
+        macaroon.add_first_party_caveat("gen = 1")
+        macaroon.add_first_party_caveat("type = access")
+
+        with self.assertRaises(AuthError) as cm:
+            yield self.auth._get_user_from_macaroon(macaroon.serialize())
+        self.assertEqual(401, cm.exception.code)
+        self.assertIn("No user caveat", cm.exception.msg)
+
+    @defer.inlineCallbacks
+    def test_get_user_from_macaroon_wrong_key(self):
+        # TODO(danielwh): Remove this mock when we remove the
+        # get_user_by_access_token fallback.
+        self.store.get_user_by_access_token = Mock(
+            return_value={"name": "@baldrick:matrix.org"}
+        )
+
+        user = "@baldrick:matrix.org"
+        macaroon = pymacaroons.Macaroon(
+            location=self.hs.config.server_name,
+            identifier="key",
+            key=self.hs.config.macaroon_secret_key + "wrong")
+        macaroon.add_first_party_caveat("gen = 1")
+        macaroon.add_first_party_caveat("type = access")
+        macaroon.add_first_party_caveat("user_id = %s" % (user,))
+
+        with self.assertRaises(AuthError) as cm:
+            yield self.auth._get_user_from_macaroon(macaroon.serialize())
+        self.assertEqual(401, cm.exception.code)
+        self.assertIn("Invalid macaroon", cm.exception.msg)
+
+    @defer.inlineCallbacks
+    def test_get_user_from_macaroon_unknown_caveat(self):
+        # TODO(danielwh): Remove this mock when we remove the
+        # get_user_by_access_token fallback.
+        self.store.get_user_by_access_token = Mock(
+            return_value={"name": "@baldrick:matrix.org"}
+        )
+
+        user = "@baldrick:matrix.org"
+        macaroon = pymacaroons.Macaroon(
+            location=self.hs.config.server_name,
+            identifier="key",
+            key=self.hs.config.macaroon_secret_key)
+        macaroon.add_first_party_caveat("gen = 1")
+        macaroon.add_first_party_caveat("type = access")
+        macaroon.add_first_party_caveat("user_id = %s" % (user,))
+        macaroon.add_first_party_caveat("cunning > fox")
+
+        with self.assertRaises(AuthError) as cm:
+            yield self.auth._get_user_from_macaroon(macaroon.serialize())
+        self.assertEqual(401, cm.exception.code)
+        self.assertIn("Invalid macaroon", cm.exception.msg)
+
+    @defer.inlineCallbacks
+    def test_get_user_from_macaroon_expired(self):
+        # TODO(danielwh): Remove this mock when we remove the
+        # get_user_by_access_token fallback.
+        self.store.get_user_by_access_token = Mock(
+            return_value={"name": "@baldrick:matrix.org"}
+        )
+
+        self.store.get_user_by_access_token = Mock(
+            return_value={"name": "@baldrick:matrix.org"}
+        )
+
+        user = "@baldrick:matrix.org"
+        macaroon = pymacaroons.Macaroon(
+            location=self.hs.config.server_name,
+            identifier="key",
+            key=self.hs.config.macaroon_secret_key)
+        macaroon.add_first_party_caveat("gen = 1")
+        macaroon.add_first_party_caveat("type = access")
+        macaroon.add_first_party_caveat("user_id = %s" % (user,))
+        macaroon.add_first_party_caveat("time < 1") # ms
+
+        self.hs.clock.now = 5000 # seconds
+
+        yield self.auth._get_user_from_macaroon(macaroon.serialize())
+        # TODO(daniel): Turn on the check that we validate expiration, when we
+        # validate expiration (and remove the above line, which will start
+        # throwing).
+        # with self.assertRaises(AuthError) as cm:
+        #     yield self.auth._get_user_from_macaroon(macaroon.serialize())
+        # self.assertEqual(401, cm.exception.code)
+        # self.assertIn("Invalid macaroon", cm.exception.msg)
diff --git a/tests/rest/client/v1/test_presence.py b/tests/rest/client/v1/test_presence.py
index 91547bdd06..29d9bbaad4 100644
--- a/tests/rest/client/v1/test_presence.py
+++ b/tests/rest/client/v1/test_presence.py
@@ -41,6 +41,22 @@ myid = "@apple:test"
 PATH_PREFIX = "/_matrix/client/api/v1"
 
 
+class NullSource(object):
+    """This event source never yields any events and its token remains at
+    zero. It may be useful for unit-testing."""
+    def __init__(self, hs):
+        pass
+
+    def get_new_events_for_user(self, user, from_key, limit):
+        return defer.succeed(([], from_key))
+
+    def get_current_key(self, direction='f'):
+        return defer.succeed(0)
+
+    def get_pagination_rows(self, user, pagination_config, key):
+        return defer.succeed(([], pagination_config.from_key))
+
+
 class JustPresenceHandlers(object):
     def __init__(self, hs):
         self.presence_handler = PresenceHandler(hs)
@@ -76,7 +92,7 @@ class PresenceStateTestCase(unittest.TestCase):
                 "token_id": 1,
             }
 
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         room_member_handler = hs.handlers.room_member_handler = Mock(
             spec=[
@@ -169,7 +185,7 @@ class PresenceListTestCase(unittest.TestCase):
             ]
         )
 
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         presence.register_servlets(hs, self.mock_resource)
 
@@ -243,7 +259,7 @@ class PresenceEventStreamTestCase(unittest.TestCase):
         # HIDEOUS HACKERY
         # TODO(paul): This should be injected in via the HomeServer DI system
         from synapse.streams.events import (
-            PresenceEventSource, NullSource, EventSources
+            PresenceEventSource, EventSources
         )
 
         old_SOURCE_TYPES = EventSources.SOURCE_TYPES
diff --git a/tests/rest/client/v1/test_rooms.py b/tests/rest/client/v1/test_rooms.py
index 34ab47d02e..a2123be81b 100644
--- a/tests/rest/client/v1/test_rooms.py
+++ b/tests/rest/client/v1/test_rooms.py
@@ -59,7 +59,7 @@ class RoomPermissionsTestCase(RestTestCase):
                 "user": UserID.from_string(self.auth_user_id),
                 "token_id": 1,
             }
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
@@ -239,7 +239,7 @@ class RoomPermissionsTestCase(RestTestCase):
                            "PUT", topic_path, topic_content)
         self.assertEquals(403, code, msg=str(response))
         (code, response) = yield self.mock_resource.trigger_get(topic_path)
-        self.assertEquals(403, code, msg=str(response))
+        self.assertEquals(200, code, msg=str(response))
 
         # get topic in PUBLIC room, not joined, expect 403
         (code, response) = yield self.mock_resource.trigger_get(
@@ -301,11 +301,11 @@ class RoomPermissionsTestCase(RestTestCase):
             room=room, expect_code=200)
 
         # get membership of self, get membership of other, private room + left
-        # expect all 403s
+        # expect all 200s
         yield self.leave(room=room, user=self.user_id)
         yield self._test_get_membership(
             members=[self.user_id, self.rmcreator_id],
-            room=room, expect_code=403)
+            room=room, expect_code=200)
 
     @defer.inlineCallbacks
     def test_membership_public_room_perms(self):
@@ -326,11 +326,11 @@ class RoomPermissionsTestCase(RestTestCase):
             room=room, expect_code=200)
 
         # get membership of self, get membership of other, public room + left
-        # expect 403.
+        # expect 200.
         yield self.leave(room=room, user=self.user_id)
         yield self._test_get_membership(
             members=[self.user_id, self.rmcreator_id],
-            room=room, expect_code=403)
+            room=room, expect_code=200)
 
     @defer.inlineCallbacks
     def test_invited_permissions(self):
@@ -444,7 +444,7 @@ class RoomsMemberListTestCase(RestTestCase):
                 "user": UserID.from_string(self.auth_user_id),
                 "token_id": 1,
             }
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
@@ -492,9 +492,9 @@ class RoomsMemberListTestCase(RestTestCase):
         self.assertEquals(200, code, msg=str(response))
 
         yield self.leave(room=room_id, user=self.user_id)
-        # can no longer see list, you've left.
+        # can see old list once left
         (code, response) = yield self.mock_resource.trigger_get(room_path)
-        self.assertEquals(403, code, msg=str(response))
+        self.assertEquals(200, code, msg=str(response))
 
 
 class RoomsCreateTestCase(RestTestCase):
@@ -522,7 +522,7 @@ class RoomsCreateTestCase(RestTestCase):
                 "user": UserID.from_string(self.auth_user_id),
                 "token_id": 1,
             }
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
@@ -614,7 +614,7 @@ class RoomTopicTestCase(RestTestCase):
                 "token_id": 1,
             }
 
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
@@ -718,7 +718,7 @@ class RoomMemberStateTestCase(RestTestCase):
                 "user": UserID.from_string(self.auth_user_id),
                 "token_id": 1,
             }
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
@@ -843,7 +843,7 @@ class RoomMessagesTestCase(RestTestCase):
                 "user": UserID.from_string(self.auth_user_id),
                 "token_id": 1,
             }
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
@@ -938,7 +938,7 @@ class RoomInitialSyncTestCase(RestTestCase):
                 "user": UserID.from_string(self.auth_user_id),
                 "token_id": 1,
             }
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
diff --git a/tests/rest/client/v1/test_typing.py b/tests/rest/client/v1/test_typing.py
index 1c4519406d..6395ce79db 100644
--- a/tests/rest/client/v1/test_typing.py
+++ b/tests/rest/client/v1/test_typing.py
@@ -67,7 +67,7 @@ class RoomTypingTestCase(RestTestCase):
                 "token_id": 1,
             }
 
-        hs.get_v1auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_v1auth()._get_user_by_access_token = _get_user_by_access_token
 
         def _insert_client_ip(*args, **kwargs):
             return defer.succeed(None)
diff --git a/tests/rest/client/v1/utils.py b/tests/rest/client/v1/utils.py
index c472d53043..85096a0326 100644
--- a/tests/rest/client/v1/utils.py
+++ b/tests/rest/client/v1/utils.py
@@ -37,9 +37,6 @@ class RestTestCase(unittest.TestCase):
         self.mock_resource = None
         self.auth_user_id = None
 
-    def mock_get_user_by_access_token(self, token=None):
-        return self.auth_user_id
-
     @defer.inlineCallbacks
     def create_room_as(self, room_creator, is_public=True, tok=None):
         temp_id = self.auth_user_id
diff --git a/tests/rest/client/v2_alpha/__init__.py b/tests/rest/client/v2_alpha/__init__.py
index ef972a53aa..f45570a1c0 100644
--- a/tests/rest/client/v2_alpha/__init__.py
+++ b/tests/rest/client/v2_alpha/__init__.py
@@ -48,7 +48,7 @@ class V2AlphaRestTestCase(unittest.TestCase):
                 "user": UserID.from_string(self.USER_ID),
                 "token_id": 1,
             }
-        hs.get_auth().get_user_by_access_token = _get_user_by_access_token
+        hs.get_auth()._get_user_by_access_token = _get_user_by_access_token
 
         for r in self.TO_REGISTER:
             r.register_servlets(hs, self.mock_resource)
diff --git a/tests/storage/event_injector.py b/tests/storage/event_injector.py
new file mode 100644
index 0000000000..42bd8928bd
--- /dev/null
+++ b/tests/storage/event_injector.py
@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+# Copyright 2015 OpenMarket 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 import unittest
+from twisted.internet import defer
+
+from synapse.api.constants import EventTypes, Membership
+from synapse.types import UserID, RoomID
+
+from tests.utils import setup_test_homeserver
+
+from mock import Mock
+
+
+class EventInjector:
+    def __init__(self, hs):
+        self.hs = hs
+        self.store = hs.get_datastore()
+        self.message_handler = hs.get_handlers().message_handler
+        self.event_builder_factory = hs.get_event_builder_factory()
+
+    @defer.inlineCallbacks
+    def create_room(self, room):
+        builder = self.event_builder_factory.new({
+            "type": EventTypes.Create,
+            "room_id": room.to_string(),
+            "content": {},
+        })
+
+        event, context = yield self.message_handler._create_new_client_event(
+            builder
+        )
+
+        yield self.store.persist_event(event, context)
+
+    @defer.inlineCallbacks
+    def inject_room_member(self, room, user, membership):
+        builder = self.event_builder_factory.new({
+            "type": EventTypes.Member,
+            "sender": user.to_string(),
+            "state_key": user.to_string(),
+            "room_id": room.to_string(),
+            "content": {"membership": membership},
+        })
+
+        event, context = yield self.message_handler._create_new_client_event(
+            builder
+        )
+
+        yield self.store.persist_event(event, context)
+
+        defer.returnValue(event)
+
+    @defer.inlineCallbacks
+    def inject_message(self, room, user, body):
+        builder = self.event_builder_factory.new({
+            "type": EventTypes.Message,
+            "sender": user.to_string(),
+            "state_key": user.to_string(),
+            "room_id": room.to_string(),
+            "content": {"body": body, "msgtype": u"message"},
+        })
+
+        event, context = yield self.message_handler._create_new_client_event(
+            builder
+        )
+
+        yield self.store.persist_event(event, context)
diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py
index 8573f18b55..1ddca1da4c 100644
--- a/tests/storage/test_base.py
+++ b/tests/storage/test_base.py
@@ -186,26 +186,6 @@ class SQLBaseStoreTestCase(unittest.TestCase):
         )
 
     @defer.inlineCallbacks
-    def test_update_one_with_return(self):
-        self.mock_txn.rowcount = 1
-        self.mock_txn.fetchone.return_value = ("Old Value",)
-
-        ret = yield self.datastore._simple_selectupdate_one(
-                table="tablename",
-                keyvalues={"keycol": "TheKey"},
-                updatevalues={"columname": "New Value"},
-                retcols=["columname"]
-        )
-
-        self.assertEquals({"columname": "Old Value"}, ret)
-        self.mock_txn.execute.assert_has_calls([
-                call('SELECT columname FROM tablename WHERE keycol = ?',
-                    ['TheKey']),
-                call("UPDATE tablename SET columname = ? WHERE keycol = ?",
-                    ["New Value", "TheKey"])
-        ])
-
-    @defer.inlineCallbacks
     def test_delete_one(self):
         self.mock_txn.rowcount = 1
 
diff --git a/tests/storage/test_events.py b/tests/storage/test_events.py
new file mode 100644
index 0000000000..313013009e
--- /dev/null
+++ b/tests/storage/test_events.py
@@ -0,0 +1,116 @@
+# -*- coding: utf-8 -*-
+# Copyright 2015 OpenMarket 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.
+import uuid
+from mock.mock import Mock
+from synapse.types import RoomID, UserID
+
+from tests import unittest
+from twisted.internet import defer
+from tests.storage.event_injector import EventInjector
+
+from tests.utils import setup_test_homeserver
+
+
+class EventsStoreTestCase(unittest.TestCase):
+
+    @defer.inlineCallbacks
+    def setUp(self):
+        self.hs = yield setup_test_homeserver(
+            resource_for_federation=Mock(),
+            http_client=None,
+        )
+        self.store = self.hs.get_datastore()
+        self.db_pool = self.hs.get_db_pool()
+        self.message_handler = self.hs.get_handlers().message_handler
+        self.event_injector = EventInjector(self.hs)
+
+    @defer.inlineCallbacks
+    def test_count_daily_messages(self):
+        self.db_pool.runQuery("DELETE FROM stats_reporting")
+
+        self.hs.clock.now = 100
+
+        # Never reported before, and nothing which could be reported
+        count = yield self.store.count_daily_messages()
+        self.assertIsNone(count)
+        count = yield self.db_pool.runQuery("SELECT COUNT(*) FROM stats_reporting")
+        self.assertEqual([(0,)], count)
+
+        # Create something to report
+        room = RoomID.from_string("!abc123:test")
+        user = UserID.from_string("@raccoonlover:test")
+        yield self.event_injector.create_room(room)
+
+        self.base_event = yield self._get_last_stream_token()
+
+        yield self.event_injector.inject_message(room, user, "Raccoons are really cute")
+
+        # Never reported before, something could be reported, but isn't because
+        # it isn't old enough.
+        count = yield self.store.count_daily_messages()
+        self.assertIsNone(count)
+        self._assert_stats_reporting(1, self.hs.clock.now)
+
+        # Already reported yesterday, two new events from today.
+        yield self.event_injector.inject_message(room, user, "Yeah they are!")
+        yield self.event_injector.inject_message(room, user, "Incredibly!")
+        self.hs.clock.now += 60 * 60 * 24
+        count = yield self.store.count_daily_messages()
+        self.assertEqual(2, count)  # 2 since yesterday
+        self._assert_stats_reporting(3, self.hs.clock.now)  # 3 ever
+
+        # Last reported too recently.
+        yield self.event_injector.inject_message(room, user, "Who could disagree?")
+        self.hs.clock.now += 60 * 60 * 22
+        count = yield self.store.count_daily_messages()
+        self.assertIsNone(count)
+        self._assert_stats_reporting(4, self.hs.clock.now)
+
+        # Last reported too long ago
+        yield self.event_injector.inject_message(room, user, "No one.")
+        self.hs.clock.now += 60 * 60 * 26
+        count = yield self.store.count_daily_messages()
+        self.assertIsNone(count)
+        self._assert_stats_reporting(5, self.hs.clock.now)
+
+        # And now let's actually report something
+        yield self.event_injector.inject_message(room, user, "Indeed.")
+        yield self.event_injector.inject_message(room, user, "Indeed.")
+        yield self.event_injector.inject_message(room, user, "Indeed.")
+        # A little over 24 hours is fine :)
+        self.hs.clock.now += (60 * 60 * 24) + 50
+        count = yield self.store.count_daily_messages()
+        self.assertEqual(3, count)
+        self._assert_stats_reporting(8, self.hs.clock.now)
+
+    @defer.inlineCallbacks
+    def _get_last_stream_token(self):
+        rows = yield self.db_pool.runQuery(
+            "SELECT stream_ordering"
+            " FROM events"
+            " ORDER BY stream_ordering DESC"
+            " LIMIT 1"
+        )
+        if not rows:
+            defer.returnValue(0)
+        else:
+            defer.returnValue(rows[0][0])
+
+    @defer.inlineCallbacks
+    def _assert_stats_reporting(self, messages, time):
+        rows = yield self.db_pool.runQuery(
+            "SELECT reported_stream_token, reported_time FROM stats_reporting"
+        )
+        self.assertEqual([(self.base_event + messages, time,)], rows)
diff --git a/tests/storage/test_room.py b/tests/storage/test_room.py
index ab7625a3ca..caffce64e3 100644
--- a/tests/storage/test_room.py
+++ b/tests/storage/test_room.py
@@ -85,7 +85,7 @@ class RoomEventsStoreTestCase(unittest.TestCase):
         # Room events need the full datastore, for persist_event() and
         # get_room_state()
         self.store = hs.get_datastore()
-        self.event_factory = hs.get_event_factory();
+        self.event_factory = hs.get_event_factory()
 
         self.room = RoomID.from_string("!abcde:test")
 
diff --git a/tests/storage/test_stream.py b/tests/storage/test_stream.py
index 0c9b89d765..a658a789aa 100644
--- a/tests/storage/test_stream.py
+++ b/tests/storage/test_stream.py
@@ -19,6 +19,7 @@ from twisted.internet import defer
 
 from synapse.api.constants import EventTypes, Membership
 from synapse.types import UserID, RoomID
+from tests.storage.event_injector import EventInjector
 
 from tests.utils import setup_test_homeserver
 
@@ -36,6 +37,7 @@ class StreamStoreTestCase(unittest.TestCase):
 
         self.store = hs.get_datastore()
         self.event_builder_factory = hs.get_event_builder_factory()
+        self.event_injector = EventInjector(hs)
         self.handlers = hs.get_handlers()
         self.message_handler = self.handlers.message_handler
 
@@ -45,60 +47,20 @@ class StreamStoreTestCase(unittest.TestCase):
         self.room1 = RoomID.from_string("!abc123:test")
         self.room2 = RoomID.from_string("!xyx987:test")
 
-        self.depth = 1
-
-    @defer.inlineCallbacks
-    def inject_room_member(self, room, user, membership):
-        self.depth += 1
-
-        builder = self.event_builder_factory.new({
-            "type": EventTypes.Member,
-            "sender": user.to_string(),
-            "state_key": user.to_string(),
-            "room_id": room.to_string(),
-            "content": {"membership": membership},
-        })
-
-        event, context = yield self.message_handler._create_new_client_event(
-            builder
-        )
-
-        yield self.store.persist_event(event, context)
-
-        defer.returnValue(event)
-
-    @defer.inlineCallbacks
-    def inject_message(self, room, user, body):
-        self.depth += 1
-
-        builder = self.event_builder_factory.new({
-            "type": EventTypes.Message,
-            "sender": user.to_string(),
-            "state_key": user.to_string(),
-            "room_id": room.to_string(),
-            "content": {"body": body, "msgtype": u"message"},
-        })
-
-        event, context = yield self.message_handler._create_new_client_event(
-            builder
-        )
-
-        yield self.store.persist_event(event, context)
-
     @defer.inlineCallbacks
     def test_event_stream_get_other(self):
         # Both bob and alice joins the room
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_alice, Membership.JOIN
         )
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_bob, Membership.JOIN
         )
 
         # Initial stream key:
         start = yield self.store.get_room_events_max_id()
 
-        yield self.inject_message(self.room1, self.u_alice, u"test")
+        yield self.event_injector.inject_message(self.room1, self.u_alice, u"test")
 
         end = yield self.store.get_room_events_max_id()
 
@@ -125,17 +87,17 @@ class StreamStoreTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def test_event_stream_get_own(self):
         # Both bob and alice joins the room
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_alice, Membership.JOIN
         )
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_bob, Membership.JOIN
         )
 
         # Initial stream key:
         start = yield self.store.get_room_events_max_id()
 
-        yield self.inject_message(self.room1, self.u_alice, u"test")
+        yield self.event_injector.inject_message(self.room1, self.u_alice, u"test")
 
         end = yield self.store.get_room_events_max_id()
 
@@ -162,22 +124,22 @@ class StreamStoreTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def test_event_stream_join_leave(self):
         # Both bob and alice joins the room
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_alice, Membership.JOIN
         )
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_bob, Membership.JOIN
         )
 
         # Then bob leaves again.
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_bob, Membership.LEAVE
         )
 
         # Initial stream key:
         start = yield self.store.get_room_events_max_id()
 
-        yield self.inject_message(self.room1, self.u_alice, u"test")
+        yield self.event_injector.inject_message(self.room1, self.u_alice, u"test")
 
         end = yield self.store.get_room_events_max_id()
 
@@ -193,17 +155,17 @@ class StreamStoreTestCase(unittest.TestCase):
 
     @defer.inlineCallbacks
     def test_event_stream_prev_content(self):
-        yield self.inject_room_member(
+        yield self.event_injector.inject_room_member(
             self.room1, self.u_bob, Membership.JOIN
         )
 
-        event1 = yield self.inject_room_member(
+        event1 = yield self.event_injector.inject_room_member(
             self.room1, self.u_alice, Membership.JOIN
         )
 
         start = yield self.store.get_room_events_max_id()
 
-        event2 = yield self.inject_room_member(
+        event2 = yield self.event_injector.inject_room_member(
             self.room1, self.u_alice, Membership.JOIN,
         )