1 files changed, 4 insertions, 6 deletions
diff --git a/tests/handlers/test_federation.py b/tests/handlers/test_federation.py
index 6fc3d8f7fd..eb6b7c22ef 100644
--- a/tests/handlers/test_federation.py
+++ b/tests/handlers/test_federation.py
@@ -14,7 +14,7 @@
 
 
 from twisted.internet import defer
-from twisted.trial import unittest
+from tests import unittest
 
 from synapse.api.events.room import (
     InviteJoinEvent, MessageEvent, RoomMemberEvent
@@ -26,12 +26,8 @@ from synapse.federation.units import Pdu
 
 from mock import NonCallableMock, ANY
 
-import logging
-
 from ..utils import get_mock_call_args
 
-logging.getLogger().addHandler(logging.NullHandler())
-
 
 class FederationTestCase(unittest.TestCase):
 
@@ -78,7 +74,9 @@ class FederationTestCase(unittest.TestCase):
 
         yield self.handlers.federation_handler.on_receive_pdu(pdu, False)
 
-        self.datastore.persist_event.assert_called_once_with(ANY, False)
+        self.datastore.persist_event.assert_called_once_with(
+            ANY, False, is_new_state=False
+        )
         self.notifier.on_new_room_event.assert_called_once_with(ANY)
 
     @defer.inlineCallbacks