diff --git a/task_sdk/tests/execution_time/test_task_runner.py b/task_sdk/tests/execution_time/test_task_runner.py
index b882aca2e9191..fcf2f3c17f6ac 100644
--- a/task_sdk/tests/execution_time/test_task_runner.py
+++ b/task_sdk/tests/execution_time/test_task_runner.py
@@ -1504,7 +1504,7 @@ def execute(self, context):
         log = mock.MagicMock()
 
         state, _ = run(runtime_ti, log)
-        finalize(ti, state, log)
+        finalize(runtime_ti, state, log)
 
         assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.SUCCESS]
 
@@ -1542,6 +1542,6 @@ def execute(self, context):
         log = mock.MagicMock()
 
         state, _ = run(runtime_ti, log)
-        finalize(ti, state, log)
+        finalize(runtime_ti, state, log)
 
         assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.FAILED]