diff --git a/tests/storage/test_event_federation.py b/tests/storage/test_event_federation.py
index 86c7ac350d..b58386994e 100644
--- a/tests/storage/test_event_federation.py
+++ b/tests/storage/test_event_federation.py
@@ -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])
|
