1 files changed, 80 insertions, 0 deletions
diff --git a/tests/util/test_async_helpers.py b/tests/util/test_async_helpers.py
index daacc54c72..9d5010bf92 100644
--- a/tests/util/test_async_helpers.py
+++ b/tests/util/test_async_helpers.py
@@ -28,6 +28,7 @@ from synapse.logging.context import (
     make_deferred_yieldable,
 )
 from synapse.util.async_helpers import (
+    AwakenableSleeper,
     ObservableDeferred,
     concurrently_execute,
     delay_cancellation,
@@ -35,6 +36,7 @@ from synapse.util.async_helpers import (
     timeout_deferred,
 )
 
+from tests.server import get_clock
 from tests.unittest import TestCase
 
 
@@ -496,3 +498,81 @@ class DelayCancellationTests(TestCase):
         # logging context.
         blocking_d.callback(None)
         self.successResultOf(d)
+
+
+class AwakenableSleeperTests(TestCase):
+    "Tests AwakenableSleeper"
+
+    def test_sleep(self):
+        reactor, _ = get_clock()
+        sleeper = AwakenableSleeper(reactor)
+
+        d = defer.ensureDeferred(sleeper.sleep("name", 1000))
+
+        reactor.pump([0.0])
+        self.assertFalse(d.called)
+
+        reactor.advance(0.5)
+        self.assertFalse(d.called)
+
+        reactor.advance(0.6)
+        self.assertTrue(d.called)
+
+    def test_explicit_wake(self):
+        reactor, _ = get_clock()
+        sleeper = AwakenableSleeper(reactor)
+
+        d = defer.ensureDeferred(sleeper.sleep("name", 1000))
+
+        reactor.pump([0.0])
+        self.assertFalse(d.called)
+
+        reactor.advance(0.5)
+        self.assertFalse(d.called)
+
+        sleeper.wake("name")
+        self.assertTrue(d.called)
+
+        reactor.advance(0.6)
+
+    def test_multiple_sleepers_timeout(self):
+        reactor, _ = get_clock()
+        sleeper = AwakenableSleeper(reactor)
+
+        d1 = defer.ensureDeferred(sleeper.sleep("name", 1000))
+
+        reactor.advance(0.6)
+        self.assertFalse(d1.called)
+
+        # Add another sleeper
+        d2 = defer.ensureDeferred(sleeper.sleep("name", 1000))
+
+        # Only the first sleep should time out now.
+        reactor.advance(0.6)
+        self.assertTrue(d1.called)
+        self.assertFalse(d2.called)
+
+        reactor.advance(0.6)
+        self.assertTrue(d2.called)
+
+    def test_multiple_sleepers_wake(self):
+        reactor, _ = get_clock()
+        sleeper = AwakenableSleeper(reactor)
+
+        d1 = defer.ensureDeferred(sleeper.sleep("name", 1000))
+
+        reactor.advance(0.5)
+        self.assertFalse(d1.called)
+
+        # Add another sleeper
+        d2 = defer.ensureDeferred(sleeper.sleep("name", 1000))
+
+        # Neither should fire yet
+        reactor.advance(0.3)
+        self.assertFalse(d1.called)
+        self.assertFalse(d2.called)
+
+        # Explicitly waking both up works
+        sleeper.wake("name")
+        self.assertTrue(d1.called)
+        self.assertTrue(d2.called)