1 files changed, 5 insertions, 2 deletions
diff --git a/tests/unittest.py b/tests/unittest.py
index 8b513bb32b..a3d39920db 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -22,6 +22,7 @@ from canonicaljson import json
 
 import twisted
 import twisted.logger
+from twisted.internet.defer import Deferred
 from twisted.trial import unittest
 
 from synapse.http.server import JsonResource
@@ -281,12 +282,14 @@ class HomeserverTestCase(TestCase):
         kwargs.update(self._hs_args)
         return setup_test_homeserver(self.addCleanup, *args, **kwargs)
 
-    def pump(self):
+    def pump(self, by=0.0):
         """
         Pump the reactor enough that Deferreds will fire.
         """
-        self.reactor.pump([0.0] * 100)
+        self.reactor.pump([by] * 100)
 
     def get_success(self, d):
+        if not isinstance(d, Deferred):
+            return d
         self.pump()
         return self.successResultOf(d)