diff --git a/tests/handlers/test_federation.py b/tests/handlers/test_federation.py
index a9d6b2bb17..e386cddb38 100644
--- a/tests/handlers/test_federation.py
+++ b/tests/handlers/test_federation.py
@@ -36,7 +36,7 @@ class FederationTestCase(unittest.TestCase):
self.mock_config.signing_key = [MockKey()]
self.state_handler = NonCallableMock(spec_set=[
- "annotate_state_groups",
+ "annotate_event_with_state",
])
self.auth = NonCallableMock(spec_set=[
@@ -85,7 +85,7 @@ class FederationTestCase(unittest.TestCase):
self.datastore.persist_event.return_value = defer.succeed(None)
self.datastore.get_room.return_value = defer.succeed(True)
- self.state_handler.annotate_state_groups.return_value = (
+ self.state_handler.annotate_event_with_state.return_value = (
defer.succeed(False)
)
@@ -95,7 +95,7 @@ class FederationTestCase(unittest.TestCase):
ANY, False, is_new_state=False
)
- self.state_handler.annotate_state_groups.assert_called_once_with(
+ self.state_handler.annotate_event_with_state.assert_called_once_with(
ANY,
old_state=None,
)
diff --git a/tests/handlers/test_room.py b/tests/handlers/test_room.py
index 55c9f6e142..ee264e5ee2 100644
--- a/tests/handlers/test_room.py
+++ b/tests/handlers/test_room.py
@@ -60,7 +60,7 @@ class RoomMemberHandlerTestCase(unittest.TestCase):
]),
auth=NonCallableMock(spec_set=["check", "add_auth_events"]),
state_handler=NonCallableMock(spec_set=[
- "annotate_state_groups",
+ "annotate_event_with_state",
]),
config=self.mock_config,
)
@@ -251,7 +251,7 @@ class RoomCreationTest(unittest.TestCase):
]),
auth=NonCallableMock(spec_set=["check", "add_auth_events"]),
state_handler=NonCallableMock(spec_set=[
- "annotate_state_groups",
+ "annotate_event_with_state",
]),
ratelimiter=NonCallableMock(spec_set=[
"send_message",
@@ -282,7 +282,7 @@ class RoomCreationTest(unittest.TestCase):
def annotate(event):
event.state_events = {}
return defer.succeed(None)
- self.state_handler.annotate_state_groups.side_effect = annotate
+ self.state_handler.annotate_event_with_state.side_effect = annotate
def hosts(room):
return defer.succeed([])
@@ -311,6 +311,6 @@ class RoomCreationTest(unittest.TestCase):
self.assertEquals(user_id, join_event.user_id)
self.assertEquals(user_id, join_event.state_key)
- self.assertTrue(self.state_handler.annotate_state_groups.called)
+ self.assertTrue(self.state_handler.annotate_event_with_state.called)
self.assertTrue(self.federation.handle_new_event.called)
|