2 files changed, 48 insertions, 0 deletions
diff --git a/tests/test_test_utils.py b/tests/test_test_utils.py
index 11c5db3cd9..dcf946fc1a 100644
--- a/tests/test_test_utils.py
+++ b/tests/test_test_utils.py
@@ -28,3 +28,25 @@ class MockClockTestCase(unittest.TestCase):
         self.clock.advance_time(20)
 
         self.assertEquals(20, self.clock.time() - start_time)
+
+    def test_later(self):
+        invoked = [0, 0]
+
+        def _cb0():
+            invoked[0] = 1
+        self.clock.call_later(10, _cb0)
+
+        def _cb1():
+            invoked[1] = 1
+        self.clock.call_later(20, _cb1)
+
+        self.assertFalse(invoked[0])
+
+        self.clock.advance_time(15)
+
+        self.assertTrue(invoked[0])
+        self.assertFalse(invoked[1])
+
+        self.clock.advance_time(5)
+
+        self.assertTrue(invoked[1])
diff --git a/tests/utils.py b/tests/utils.py
index d8be73dba8..72843714fd 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -18,6 +18,8 @@ from synapse.api.errors import cs_error, CodeMessageException, StoreError
 from synapse.api.constants import Membership
 from synapse.storage import prepare_database
 
+from synapse.util.logcontext import LoggingContext
+
 from synapse.api.events.room import (
     RoomMemberEvent, MessageEvent
 )
@@ -134,16 +136,40 @@ class MockKey(object):
 class MockClock(object):
     now = 1000
 
+    def __init__(self):
+        # list of tuples of (absolute_time, callback) in no particular order
+        self.timers = []
+
     def time(self):
         return self.now
 
     def time_msec(self):
         return self.time() * 1000
 
+    def call_later(self, delay, callback):
+        current_context = LoggingContext.current_context()
+
+        def wrapped_callback():
+            LoggingContext.thread_local.current_context = current_context
+            callback()
+        self.timers.append((self.now + delay, wrapped_callback))
+
+    def cancel_call_later(self, timer):
+        raise NotImplementedError("Oopsie")
+
     # For unit testing
     def advance_time(self, secs):
         self.now += secs
 
+        timers = self.timers
+        self.timers = []
+
+        for time, callback in timers:
+            if self.now >= time:
+                callback()
+            else:
+                self.timers.append((time, callback))
+
 
 class SQLiteMemoryDbPool(ConnectionPool, object):
     def __init__(self):