diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/federation/test_federation_sender.py | 14 | ||||
| -rw-r--r-- | tests/handlers/test_typing.py | 6 |
2 files changed, 11 insertions, 9 deletions
diff --git a/tests/federation/test_federation_sender.py b/tests/federation/test_federation_sender.py index b5be727fe4..01a1db6115 100644 --- a/tests/federation/test_federation_sender.py +++ b/tests/federation/test_federation_sender.py @@ -30,16 +30,16 @@ from tests.unittest import HomeserverTestCase, override_config class FederationSenderReceiptsTestCases(HomeserverTestCase): def make_homeserver(self, reactor, clock): - mock_state_handler = Mock(spec=["get_current_hosts_in_room"]) - # Ensure a new Awaitable is created for each call. - mock_state_handler.get_current_hosts_in_room.return_value = make_awaitable( - ["test", "host2"] - ) - return self.setup_test_homeserver( - state_handler=mock_state_handler, + hs = self.setup_test_homeserver( federation_transport_client=Mock(spec=["send_transaction"]), ) + hs.get_storage_controllers().state.get_current_hosts_in_room = Mock( + return_value=make_awaitable({"test", "host2"}) + ) + + return hs + @override_config({"send_federation": True}) def test_send_receipts(self): mock_send_transaction = ( diff --git a/tests/handlers/test_typing.py b/tests/handlers/test_typing.py index 14a0ee4922..7af1333126 100644 --- a/tests/handlers/test_typing.py +++ b/tests/handlers/test_typing.py @@ -129,10 +129,12 @@ class TypingNotificationsTestCase(unittest.HomeserverTestCase): hs.get_event_auth_handler().check_host_in_room = check_host_in_room - def get_joined_hosts_for_room(room_id: str): + async def get_current_hosts_in_room(room_id: str): return {member.domain for member in self.room_members} - self.datastore.get_joined_hosts_for_room = get_joined_hosts_for_room + hs.get_storage_controllers().state.get_current_hosts_in_room = ( + get_current_hosts_in_room + ) async def get_users_in_room(room_id: str): return {str(u) for u in self.room_members} |