fixed test

Author: Andrei Neagu

packages/service-library/tests/deferred_tasks/test__base_deferred_handler.py

@@ -43,7 +43,8 @@ class MockKeys(StrAutoEnum):
     GET_TIMEOUT = auto()
     START_DEFERRED = auto()
     ON_DEFERRED_CREATED = auto()
-    RUN_DEFERRED = auto()
+    RUN_DEFERRED_BEFORE_HADNLER = auto()
+    RUN_DEFERRED_AFTER_HANDLER = auto()
     ON_DEFERRED_RESULT = auto()
     ON_FINISHED_WITH_ERROR = auto()
 
@@ -123,8 +124,9 @@ async def on_created(
 
             @classmethod
             async def run(cls, context: DeferredContext) -> Any:
+                mocks[MockKeys.RUN_DEFERRED_BEFORE_HADNLER](context)
                 result = await run(context)
-                mocks[MockKeys.RUN_DEFERRED](context)
+                mocks[MockKeys.RUN_DEFERRED_AFTER_HANDLER](context)
                 return result
 
             @classmethod
@@ -230,8 +232,8 @@ async def _run_ok(_: DeferredContext) -> Any:
     await _assert_mock_call(mocks, key=MockKeys.ON_DEFERRED_CREATED, count=1)
     assert TaskUID(mocks[MockKeys.ON_DEFERRED_CREATED].call_args_list[0].args[0])
 
-    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED, count=1)
-    mocks[MockKeys.RUN_DEFERRED].assert_called_once_with(context)
+    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED_AFTER_HANDLER, count=1)
+    mocks[MockKeys.RUN_DEFERRED_AFTER_HANDLER].assert_called_once_with(context)
 
     await _assert_mock_call(mocks, key=MockKeys.ON_DEFERRED_RESULT, count=1)
     mocks[MockKeys.ON_DEFERRED_RESULT].assert_called_once_with(run_return, context)
@@ -283,7 +285,7 @@ async def _run_raises(_: DeferredContext) -> None:
             count=retry_count,
         )
 
-    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED, count=0)
+    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED_AFTER_HANDLER, count=0)
     await _assert_mock_call(mocks, key=MockKeys.ON_DEFERRED_RESULT, count=0)
 
     await _assert_log_message(
@@ -320,6 +322,7 @@ async def _run_to_cancel(_: DeferredContext) -> None:
     await _assert_mock_call(mocks, key=MockKeys.ON_DEFERRED_CREATED, count=1)
     task_uid = TaskUID(mocks[MockKeys.ON_DEFERRED_CREATED].call_args_list[0].args[0])
 
+    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED_BEFORE_HADNLER, count=1)
     await mocked_deferred_handler.cancel(task_uid)
 
     await _assert_mock_call(mocks, key=MockKeys.ON_FINISHED_WITH_ERROR, count=0)
@@ -331,7 +334,7 @@ async def _run_to_cancel(_: DeferredContext) -> None:
         == 0
     )
 
-    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED, count=0)
+    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED_AFTER_HANDLER, count=0)
     await _assert_mock_call(mocks, key=MockKeys.ON_DEFERRED_RESULT, count=0)
 
     await _assert_log_message(
@@ -451,7 +454,7 @@ async def _run_that_times_out(_: DeferredContext) -> None:
     for entry in mocks[MockKeys.ON_FINISHED_WITH_ERROR].call_args_list:
         assert "builtins.TimeoutError" in entry.args[0].error
 
-    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED, count=0)
+    await _assert_mock_call(mocks, key=MockKeys.RUN_DEFERRED_AFTER_HANDLER, count=0)
     await _assert_mock_call(mocks, key=MockKeys.ON_DEFERRED_RESULT, count=0)