omkar-foss commented on code in PR #44220:
URL: https://github.com/apache/airflow/pull/44220#discussion_r1851896645
##########
tests/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -1663,3 +1668,649 @@ def test_raises_404_for_nonexistent_task_instance(self,
test_client, session):
assert response.json() == {
"detail": "The Task Instance with dag_id:
`example_python_operator`, run_id: `TEST_DAG_RUN_ID`, task_id:
`nonexistent_task`, try_number: `0` and map_index: `-1` was not found"
}
+
+
+class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
+ @pytest.mark.parametrize(
+ "main_dag, task_instances, request_dag, payload, expected_ti",
+ [
+ pytest.param(
+ "example_python_operator",
+ [
+ {"logical_date": DEFAULT_DATETIME_1, "state":
State.FAILED},
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=1),
+ "state": State.FAILED,
+ },
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=2),
+ "state": State.FAILED,
+ },
+ ],
+ "example_python_operator",
+ {
+ "dry_run": True,
+ "start_date": DEFAULT_DATETIME_STR_2,
+ "only_failed": True,
+ },
+ 2,
+ id="clear start date filter",
+ ),
+ pytest.param(
+ "example_python_operator",
+ [
+ {"logical_date": DEFAULT_DATETIME_1, "state":
State.FAILED},
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=1),
+ "state": State.FAILED,
+ },
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=2),
+ "state": State.FAILED,
+ },
+ ],
+ "example_python_operator",
+ {
+ "dry_run": True,
+ "end_date": DEFAULT_DATETIME_STR_2,
+ "only_failed": True,
+ },
+ 2,
+ id="clear end date filter",
+ ),
+ pytest.param(
+ "example_python_operator",
+ [
+ {"logical_date": DEFAULT_DATETIME_1, "state":
State.RUNNING},
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=1),
+ "state": State.RUNNING,
+ },
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=2),
+ "state": State.FAILED,
+ },
+ ],
+ "example_python_operator",
+ {"dry_run": True, "only_running": True, "only_failed": False},
+ 2,
+ id="clear only running",
+ ),
+ pytest.param(
+ "example_python_operator",
+ [
+ {"logical_date": DEFAULT_DATETIME_1, "state":
State.FAILED},
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=1),
+ "state": State.FAILED,
+ },
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=2),
+ "state": State.RUNNING,
+ },
+ ],
+ "example_python_operator",
+ {
+ "dry_run": True,
+ "only_failed": True,
+ },
+ 2,
+ id="clear only failed",
+ ),
+ pytest.param(
+ "example_python_operator",
+ [
+ {"logical_date": DEFAULT_DATETIME_1, "state":
State.FAILED},
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=1),
+ "state": State.FAILED,
+ },
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=2),
+ "state": State.FAILED,
+ },
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=3),
+ "state": State.FAILED,
+ },
+ ],
+ "example_python_operator",
+ {
+ "dry_run": True,
+ "task_ids": ["print_the_context", "sleep_for_1"],
+ },
+ 2,
+ id="clear by task ids",
+ ),
+ pytest.param(
+ "example_python_operator",
+ [
+ {"logical_date": DEFAULT_DATETIME_1, "state":
State.FAILED},
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=1),
+ "state": State.FAILED,
+ },
+ {
+ "logical_date": DEFAULT_DATETIME_1 +
dt.timedelta(days=2),
+ "state": State.RUNNING,
+ },
+ ],
+ "example_python_operator",
+ {
+ "only_failed": True,
+ },
+ 2,
+ id="dry_run default",
+ ),
+ ],
+ )
+ def test_should_respond_200(
+ self,
+ test_client,
+ session,
+ main_dag,
+ task_instances,
+ request_dag,
+ payload,
+ expected_ti,
+ ):
+ self.create_task_instances(
+ session,
+ dag_id=main_dag,
+ task_instances=task_instances,
+ update_extras=False,
+ )
+ self.dagbag.sync_to_db()
+ response = test_client.post(
+ f"/public/dags/{request_dag}/clearTaskInstances",
+ json=payload,
+ )
+ assert response.status_code == 200
+ assert len(response.json()["task_instances"]) == expected_ti
+
+ def test_clear_taskinstance_is_called_with_queued_dr_state(self,
test_client, session):
+ """Test that if reset_dag_runs is True, then clear_task_instances is
called with State.QUEUED"""
+ self.create_task_instances(session)
+ dag_id = "example_python_operator"
+ payload = {"reset_dag_runs": True, "dry_run": False}
+ self.dagbag.sync_to_db()
+ with mock.patch(
+
"airflow.api_fastapi.core_api.routes.public.task_instances.clear_task_instances",
+ ) as mp:
Review Comment:
Done, added the decorator back from legacy.
--
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]
