diff --git a/tests/appservice/test_scheduler.py b/tests/appservice/test_scheduler.py
index a2b5ed2030..55f0899bae 100644
--- a/tests/appservice/test_scheduler.py
+++ b/tests/appservice/test_scheduler.py
@@ -24,7 +24,7 @@ from synapse.appservice.scheduler import (
from synapse.logging.context import make_deferred_yieldable
from tests import unittest
-from tests.test_utils import make_awaitable
+from tests.test_utils import simple_async_mock
from ..utils import MockClock
@@ -49,11 +49,10 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
txn = Mock(id=txn_id, service=service, events=events)
# mock methods
- self.store.get_appservice_state = Mock(
- return_value=defer.succeed(ApplicationServiceState.UP)
- )
- txn.send = Mock(return_value=make_awaitable(True))
- self.store.create_appservice_txn = Mock(return_value=defer.succeed(txn))
+ self.store.get_appservice_state = simple_async_mock(ApplicationServiceState.UP)
+ txn.send = simple_async_mock(True)
+ txn.complete = simple_async_mock(True)
+ self.store.create_appservice_txn = simple_async_mock(txn)
# actual call
self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
@@ -71,10 +70,10 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
events = [Mock(), Mock()]
txn = Mock(id="idhere", service=service, events=events)
- self.store.get_appservice_state = Mock(
- return_value=defer.succeed(ApplicationServiceState.DOWN)
+ self.store.get_appservice_state = simple_async_mock(
+ ApplicationServiceState.DOWN
)
- self.store.create_appservice_txn = Mock(return_value=defer.succeed(txn))
+ self.store.create_appservice_txn = simple_async_mock(txn)
# actual call
self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
@@ -94,12 +93,10 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
txn = Mock(id=txn_id, service=service, events=events)
# mock methods
- self.store.get_appservice_state = Mock(
- return_value=defer.succeed(ApplicationServiceState.UP)
- )
- self.store.set_appservice_state = Mock(return_value=defer.succeed(True))
- txn.send = Mock(return_value=make_awaitable(False)) # fails to send
- self.store.create_appservice_txn = Mock(return_value=defer.succeed(txn))
+ self.store.get_appservice_state = simple_async_mock(ApplicationServiceState.UP)
+ self.store.set_appservice_state = simple_async_mock(True)
+ txn.send = simple_async_mock(False) # fails to send
+ self.store.create_appservice_txn = simple_async_mock(txn)
# actual call
self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
@@ -122,7 +119,7 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.as_api = Mock()
self.store = Mock()
self.service = Mock()
- self.callback = Mock()
+ self.callback = simple_async_mock()
self.recoverer = _Recoverer(
clock=self.clock,
as_api=self.as_api,
@@ -144,8 +141,8 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.recoverer.recover()
# shouldn't have called anything prior to waiting for exp backoff
self.assertEquals(0, self.store.get_oldest_unsent_txn.call_count)
- txn.send = Mock(return_value=make_awaitable(True))
- txn.complete.return_value = make_awaitable(None)
+ txn.send = simple_async_mock(True)
+ txn.complete = simple_async_mock(None)
# wait for exp backoff
self.clock.advance_time(2)
self.assertEquals(1, txn.send.call_count)
@@ -170,8 +167,8 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.recoverer.recover()
self.assertEquals(0, self.store.get_oldest_unsent_txn.call_count)
- txn.send = Mock(return_value=make_awaitable(False))
- txn.complete.return_value = make_awaitable(None)
+ txn.send = simple_async_mock(False)
+ txn.complete = simple_async_mock(None)
self.clock.advance_time(2)
self.assertEquals(1, txn.send.call_count)
self.assertEquals(0, txn.complete.call_count)
@@ -184,7 +181,7 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.assertEquals(3, txn.send.call_count)
self.assertEquals(0, txn.complete.call_count)
self.assertEquals(0, self.callback.call_count)
- txn.send = Mock(return_value=make_awaitable(True)) # successfully send the txn
+ txn.send = simple_async_mock(True) # successfully send the txn
pop_txn = True # returns the txn the first time, then no more.
self.clock.advance_time(16)
self.assertEquals(1, txn.send.call_count) # new mock reset call count
@@ -195,6 +192,7 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
class ApplicationServiceSchedulerQueuerTestCase(unittest.TestCase):
def setUp(self):
self.txn_ctrl = Mock()
+ self.txn_ctrl.send = simple_async_mock()
self.queuer = _ServiceQueuer(self.txn_ctrl, MockClock())
def test_send_single_event_no_queue(self):
|