Add a replication stream for state groups

1 files changed, 27 insertions, 3 deletions

diff --git a/tests/replication/test_resource.py b/tests/replication/test_resource.py
index f4b5fb3328..b1dd7b4a74 100644
--- a/tests/replication/test_resource.py
+++ b/tests/replication/test_resource.py
@@ -58,15 +58,21 @@ class ReplicationResourceCase(unittest.TestCase):
         self.assertEquals(body, {})
 
     @defer.inlineCallbacks
-    def test_events(self):
-        get = self.get(events="-1", timeout="0")
+    def test_events_and_state(self):
+        get = self.get(events="-1", state="-1", timeout="0")
         yield self.hs.get_handlers().room_creation_handler.create_room(
             Requester(self.user, "", False), {}
         )
         code, body = yield get
         self.assertEquals(code, 200)
         self.assertEquals(body["events"]["field_names"], [
-            "position", "internal", "json"
+            "position", "internal", "json", "state_group"
+        ])
+        self.assertEquals(body["state_groups"]["field_names"], [
+            "position", "room_id", "event_id"
+        ])
+        self.assertEquals(body["state_group_state"]["field_names"], [
+            "position", "type", "state_key", "event_id"
         ])
 
     @defer.inlineCallbacks
@@ -132,6 +138,7 @@ class ReplicationResourceCase(unittest.TestCase):
     test_timeout_backfill = _test_timeout("backfill")
     test_timeout_push_rules = _test_timeout("push_rules")
     test_timeout_pushers = _test_timeout("pushers")
+    test_timeout_state = _test_timeout("state")
 
     @defer.inlineCallbacks
     def send_text_message(self, room_id, message):
@@ -182,4 +189,21 @@ class ReplicationResourceCase(unittest.TestCase):
         )
         response_body = json.loads(response_json)
 
+        if response_code == 200:
+            self.check_response(response_body)
+
         defer.returnValue((response_code, response_body))
+
+    def check_response(self, response_body):
+        for name, stream in response_body.items():
+            self.assertIn("field_names", stream)
+            field_names = stream["field_names"]
+            self.assertIn("rows", stream)
+            self.assertTrue(stream["rows"])
+            for row in stream["rows"]:
+                self.assertEquals(
+                    len(row), len(field_names),
+                    "%s: len(row = %r) == len(field_names = %r)" % (
+                        name, row, field_names
+                    )
+                )