diff --git a/tests/storage/test_devices.py b/tests/storage/test_devices.py
index d1227dd4ac..5491fbf6da 100644
--- a/tests/storage/test_devices.py
+++ b/tests/storage/test_devices.py
@@ -21,6 +21,29 @@ class DeviceStoreTestCase(HomeserverTestCase):
def prepare(self, reactor, clock, hs):
self.store = hs.get_datastores().main
+ def add_device_change(self, user_id, device_ids, host):
+ """Add a device list change for the given device to
+ `device_lists_outbound_pokes` table.
+ """
+
+ for device_id in device_ids:
+ stream_id = self.get_success(
+ self.store.add_device_change_to_streams(
+ "user_id", [device_id], ["!some:room"]
+ )
+ )
+
+ self.get_success(
+ self.store.add_device_list_outbound_pokes(
+ user_id=user_id,
+ device_id=device_id,
+ room_id="!some:room",
+ stream_id=stream_id,
+ hosts=[host],
+ context={},
+ )
+ )
+
def test_store_new_device(self):
self.get_success(
self.store.store_device("user_id", "device_id", "display_name")
@@ -95,11 +118,7 @@ class DeviceStoreTestCase(HomeserverTestCase):
device_ids = ["device_id1", "device_id2"]
# Add two device updates with sequential `stream_id`s
- self.get_success(
- self.store.add_device_change_to_streams(
- "user_id", device_ids, ["somehost"], ["!some:room"]
- )
- )
+ self.add_device_change("user_id", device_ids, "somehost")
# Get all device updates ever meant for this remote
now_stream_id, device_updates = self.get_success(
@@ -123,11 +142,7 @@ class DeviceStoreTestCase(HomeserverTestCase):
"device_id4",
"device_id5",
]
- self.get_success(
- self.store.add_device_change_to_streams(
- "user_id", device_ids, ["somehost"], ["!some:room"]
- )
- )
+ self.add_device_change("user_id", device_ids, "somehost")
# Get device updates meant for this remote
next_stream_id, device_updates = self.get_success(
@@ -147,11 +162,7 @@ class DeviceStoreTestCase(HomeserverTestCase):
# Add some more device updates to ensure it still resumes properly
device_ids = ["device_id6", "device_id7"]
- self.get_success(
- self.store.add_device_change_to_streams(
- "user_id", device_ids, ["somehost"], ["!some:room"]
- )
- )
+ self.add_device_change("user_id", device_ids, "somehost")
# Get the next batch of device updates
next_stream_id, device_updates = self.get_success(
@@ -224,11 +235,7 @@ class DeviceStoreTestCase(HomeserverTestCase):
"fakeSelfSigning",
]
- self.get_success(
- self.store.add_device_change_to_streams(
- "@user_id:test", device_ids, ["somehost"], ["!some:room"]
- )
- )
+ self.add_device_change("@user_id:test", device_ids, "somehost")
# Get device updates meant for this remote
next_stream_id, device_updates = self.get_success(
|
