Fix unit tests

Author: erikjohnston

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 @@ def test_get_device_updates_by_remote(self):
         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 @@ def test_get_device_updates_by_remote_can_limit_properly(self):
             "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 @@ def test_get_device_updates_by_remote_can_limit_properly(self):
 
         # 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 @@ def test_get_device_updates_by_remote_cross_signing_key_updates(
             "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(