From b589847fe4d0cfd062a5b13f74570f4853c4dc2c Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 9 Jan 2025 21:41:41 +0530 Subject: [PATCH] Apply suggestions from code review Co-authored-by: Kaxil Naik --- task_sdk/tests/execution_time/test_task_runner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/task_sdk/tests/execution_time/test_task_runner.py b/task_sdk/tests/execution_time/test_task_runner.py index d3754cc3771b2..c16517c843038 100644 --- a/task_sdk/tests/execution_time/test_task_runner.py +++ b/task_sdk/tests/execution_time/test_task_runner.py @@ -559,7 +559,7 @@ def execute(self, context): {"push_task1", "push_task2"}, ], ) -def test_xcom_pull_behavior(mocked_parse, make_ti_context, mock_supervisor_comms, spy_agency, task_ids): +def test_xcom_pull_behavior(create_runtime_ti, mock_supervisor_comms, spy_agency, task_ids): """Test that a task pulls the expected XCom value if it exists.""" class CustomOperator(BaseOperator): @@ -573,7 +573,7 @@ def execute(self, context): ) what = StartupDetails(ti=ti, file="", requests_fd=0, ti_context=make_ti_context()) - runtime_ti = mocked_parse(what, ti.dag_id, task) + runtime_ti = create_runtime_ti(task=task) mock_supervisor_comms.xcom_pull.return_value = "xcom_value"