tests

tests/test_task_runners.py

@@ -203,7 +203,7 @@ async def test_wait_captures_exceptions_as_crashed_state(
         lack of re-raise here than the equality of the exception.
         """
 
-        async def fake_orchestrate_task_run():
+        async def fake_orchestrate_task_run(task_run):
             raise exception
 
         task_run = TaskRun(flow_run_id=uuid4(), task_key=uuid4(), dynamic_key="bar")
@@ -333,7 +333,7 @@ async def test_submit_and_wait(self, task_runner):
 
         task_run = TaskRun(flow_run_id=uuid4(), task_key="foo", dynamic_key="bar")
 
-        async def fake_orchestrate_task_run(example_kwarg):
+        async def fake_orchestrate_task_run(example_kwarg, task_run):
             return State(
                 type=StateType.COMPLETED,
                 data=example_kwarg,