amoghrajesh commented on code in PR #45106:
URL: https://github.com/apache/airflow/pull/45106#discussion_r1895924293
##########
tests/api_fastapi/execution_api/routes/test_task_instances.py:
##########
@@ -340,6 +340,71 @@ def test_ti_update_state_to_reschedule(self, client,
session, create_task_instan
assert trs[0].map_index == -1
assert trs[0].duration == 129600
+ @pytest.mark.parametrize(
+ ("retries", "expected_state"),
+ [
+ # retries given
+ (2, State.UP_FOR_RETRY),
+ # retries not given
+ (None, State.FAILED),
+ # retries given but as 0
+ (0, State.FAILED),
+ # retries not known, given as -1, calculates on table default
+ (-1, State.UP_FOR_RETRY),
+ ],
+ )
+ def test_ti_update_state_to_retry(self, client, session,
create_task_instance, retries, expected_state):
+ ti = create_task_instance(
+ task_id="test_ti_update_state_to_retry",
+ state=State.RUNNING,
+ )
+ ti.retries = retries
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json={
+ "state": State.FAILED,
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "task_retries": retries,
+ },
+ )
+
+ assert response.status_code == 204
+ assert response.text == ""
+
+ session.expire_all()
+
+ ti = session.get(TaskInstance, ti.id)
+ assert ti.state == expected_state
+ assert ti.next_method is None
+ assert ti.next_kwargs is None
+
+ def test_ti_update_state_to_retry_when_restarting(self, client, session,
create_task_instance):
+ ti = create_task_instance(
+ task_id="test_ti_update_state_to_retry_when_restarting",
+ state=State.RESTARTING,
+ )
+ session.commit()
Review Comment:
The `create_task_instance`creates a ti and returns it, which we commit right
after, no need to pass `session`, all the tests in the file have same pattern
and here is the `ti` with the state being set, running on debug mode

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]