1 files changed, 7 insertions, 8 deletions
diff --git a/tests/federation/test_federation.py b/tests/federation/test_federation.py
index 0b105fe723..bb17e9aafe 100644
--- a/tests/federation/test_federation.py
+++ b/tests/federation/test_federation.py
@@ -14,11 +14,10 @@
 
 # trial imports
 from twisted.internet import defer
-from twisted.trial import unittest
+from tests import unittest
 
 # python imports
-from mock import Mock
-import logging
+from mock import Mock, ANY
 
 from ..utils import MockHttpResource, MockClock
 
@@ -28,9 +27,6 @@ from synapse.federation.units import Pdu
 from synapse.storage.pdu import PduTuple, PduEntry
 
 
-logging.getLogger().addHandler(logging.NullHandler())
-
-
 def make_pdu(prev_pdus=[], **kwargs):
     """Provide some default fields for making a PduTuple."""
     pdu_fields = {
@@ -185,7 +181,8 @@ class FederationTestCase(unittest.TestCase):
                             "depth": 1,
                         },
                     ]
-                }
+                },
+                on_send_callback=ANY,
         )
 
     @defer.inlineCallbacks
@@ -216,7 +213,9 @@ class FederationTestCase(unittest.TestCase):
                             "content": {"testing": "content here"},
                         }
                     ],
-                })
+                },
+                on_send_callback=ANY,
+        )
 
     @defer.inlineCallbacks
     def test_recv_edu(self):