1 files changed, 41 insertions, 1 deletions
diff --git a/tests/storage/test_event_federation.py b/tests/storage/test_event_federation.py
index 86c7ac350d..2fe50377f8 100644
--- a/tests/storage/test_event_federation.py
+++ b/tests/storage/test_event_federation.py
@@ -57,7 +57,7 @@ class EventFederationWorkerStoreTestCase(tests.unittest.TestCase):
                     "(event_id, algorithm, hash) "
                     "VALUES (?, 'sha256', ?)"
                 ),
-                (event_id, b"ffff"),
+                (event_id, bytearray(b"ffff")),
             )
 
         for i in range(0, 11):
@@ -75,3 +75,43 @@ class EventFederationWorkerStoreTestCase(tests.unittest.TestCase):
             el = r[i]
             depth = el[2]
             self.assertLessEqual(5, depth)
+
+    @defer.inlineCallbacks
+    def test_get_rooms_with_many_extremities(self):
+        room1 = "#room1"
+        room2 = "#room2"
+        room3 = "#room3"
+
+        def insert_event(txn, i, room_id):
+            event_id = "$event_%i:local" % i
+            txn.execute(
+                (
+                    "INSERT INTO event_forward_extremities (room_id, event_id) "
+                    "VALUES (?, ?)"
+                ),
+                (room_id, event_id),
+            )
+
+        for i in range(0, 20):
+            yield self.store.runInteraction("insert", insert_event, i, room1)
+            yield self.store.runInteraction("insert", insert_event, i, room2)
+            yield self.store.runInteraction("insert", insert_event, i, room3)
+
+        # Test simple case
+        r = yield self.store.get_rooms_with_many_extremities(5, 5, [])
+        self.assertEqual(len(r), 3)
+
+        # Does filter work?
+
+        r = yield self.store.get_rooms_with_many_extremities(5, 5, [room1])
+        self.assertTrue(room2 in r)
+        self.assertTrue(room3 in r)
+        self.assertEqual(len(r), 2)
+
+        r = yield self.store.get_rooms_with_many_extremities(5, 5, [room1, room2])
+        self.assertEqual(r, [room3])
+
+        # Does filter and limit work?
+
+        r = yield self.store.get_rooms_with_many_extremities(5, 1, [room1])
+        self.assertTrue(r == [room2] or r == [room3])