diff --git a/tests/appservice/test_scheduler.py b/tests/appservice/test_scheduler.py
index e2a3bad065..445919417e 100644
--- a/tests/appservice/test_scheduler.py
+++ b/tests/appservice/test_scheduler.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import List, Optional, Sequence, Tuple, cast
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
from typing_extensions import TypeAlias
@@ -37,7 +37,6 @@ from synapse.types import DeviceListUpdates, JsonDict
from synapse.util import Clock
from tests import unittest
-from tests.test_utils import simple_async_mock
from ..utils import MockClock
@@ -62,10 +61,12 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
txn = Mock(id=txn_id, service=service, events=events)
# mock methods
- 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)
+ self.store.get_appservice_state = AsyncMock(
+ return_value=ApplicationServiceState.UP
+ )
+ txn.send = AsyncMock(return_value=True)
+ txn.complete = AsyncMock(return_value=True)
+ self.store.create_appservice_txn = AsyncMock(return_value=txn)
# actual call
self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
@@ -89,10 +90,10 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
events = [Mock(), Mock()]
txn = Mock(id="idhere", service=service, events=events)
- self.store.get_appservice_state = simple_async_mock(
- ApplicationServiceState.DOWN
+ self.store.get_appservice_state = AsyncMock(
+ return_value=ApplicationServiceState.DOWN
)
- self.store.create_appservice_txn = simple_async_mock(txn)
+ self.store.create_appservice_txn = AsyncMock(return_value=txn)
# actual call
self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
@@ -118,10 +119,12 @@ class ApplicationServiceSchedulerTransactionCtrlTestCase(unittest.TestCase):
txn = Mock(id=txn_id, service=service, events=events)
# mock methods
- 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)
+ self.store.get_appservice_state = AsyncMock(
+ return_value=ApplicationServiceState.UP
+ )
+ self.store.set_appservice_state = AsyncMock(return_value=True)
+ txn.send = AsyncMock(return_value=False) # fails to send
+ self.store.create_appservice_txn = AsyncMock(return_value=txn)
# actual call
self.successResultOf(defer.ensureDeferred(self.txnctrl.send(service, events)))
@@ -150,7 +153,7 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.as_api = Mock()
self.store = Mock()
self.service = Mock()
- self.callback = simple_async_mock()
+ self.callback = AsyncMock()
self.recoverer = _Recoverer(
clock=cast(Clock, self.clock),
as_api=self.as_api,
@@ -174,8 +177,8 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.recoverer.recover()
# shouldn't have called anything prior to waiting for exp backoff
self.assertEqual(0, self.store.get_oldest_unsent_txn.call_count)
- txn.send = simple_async_mock(True)
- txn.complete = simple_async_mock(None)
+ txn.send = AsyncMock(return_value=True)
+ txn.complete = AsyncMock(return_value=None)
# wait for exp backoff
self.clock.advance_time(2)
self.assertEqual(1, txn.send.call_count)
@@ -202,8 +205,8 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.recoverer.recover()
self.assertEqual(0, self.store.get_oldest_unsent_txn.call_count)
- txn.send = simple_async_mock(False)
- txn.complete = simple_async_mock(None)
+ txn.send = AsyncMock(return_value=False)
+ txn.complete = AsyncMock(return_value=None)
self.clock.advance_time(2)
self.assertEqual(1, txn.send.call_count)
self.assertEqual(0, txn.complete.call_count)
@@ -216,7 +219,7 @@ class ApplicationServiceSchedulerRecovererTestCase(unittest.TestCase):
self.assertEqual(3, txn.send.call_count)
self.assertEqual(0, txn.complete.call_count)
self.assertEqual(0, self.callback.call_count)
- txn.send = simple_async_mock(True) # successfully send the txn
+ txn.send = AsyncMock(return_value=True) # successfully send the txn
pop_txn = True # returns the txn the first time, then no more.
self.clock.advance_time(16)
self.assertEqual(1, txn.send.call_count) # new mock reset call count
@@ -244,7 +247,7 @@ class ApplicationServiceSchedulerQueuerTestCase(unittest.HomeserverTestCase):
def prepare(self, reactor: "MemoryReactor", clock: Clock, hs: HomeServer) -> None:
self.scheduler = ApplicationServiceScheduler(hs)
self.txn_ctrl = Mock()
- self.txn_ctrl.send = simple_async_mock()
+ self.txn_ctrl.send = AsyncMock()
# Replace instantiated _TransactionController instances with our Mock
self.scheduler.txn_ctrl = self.txn_ctrl
|
