summary refs log tree commit diff
path: root/tests/appservice
diff options
context:
space:
mode:
Diffstat (limited to 'tests/appservice')
-rw-r--r--tests/appservice/test_scheduler.py40
1 files changed, 19 insertions, 21 deletions
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):