diff --git a/tests/util/test_task_scheduler.py b/tests/util/test_task_scheduler.py
index e97f0ed611..f8f3c9dafd 100644
--- a/tests/util/test_task_scheduler.py
+++ b/tests/util/test_task_scheduler.py
@@ -23,6 +23,7 @@ from typing import List, Optional, Tuple
 from twisted.internet.task import deferLater
 from twisted.internet.testing import MemoryReactor
 
+from synapse.logging.context import make_deferred_yieldable
 from synapse.server import HomeServer
 from synapse.types import JsonMapping, ScheduledTask, TaskStatus
 from synapse.util.clock import Clock
@@ -87,7 +88,7 @@ class TestTaskScheduler(HomeserverTestCase):
         self, task: ScheduledTask
     ) -> Tuple[TaskStatus, Optional[JsonMapping], Optional[str]]:
         # Sleep for a second
-        await deferLater(self.reactor, 1, lambda: None)
+        await make_deferred_yieldable(deferLater(self.reactor, 1, lambda: None))
         return TaskStatus.COMPLETE, None, None
 
     def test_schedule_lot_of_tasks(self) -> None:
@@ -171,7 +172,7 @@ class TestTaskScheduler(HomeserverTestCase):
         else:
             await self.task_scheduler.update_task(task.id, result={"in_progress": True})
             # Await forever to simulate an aborted task because of a restart
-            await deferLater(self.reactor, 2**16, lambda: None)
+            await make_deferred_yieldable(deferLater(self.reactor, 2**16, lambda: None))
             # This should never been called
             return TaskStatus.ACTIVE, None, None