amoghrajesh commented on code in PR #72100:
URL: https://github.com/apache/airflow/pull/72100#discussion_r3976316997
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -4339,6 +4339,72 @@ def test_clear_dry_run_does_not_set_note(self,
test_client, session):
ti_id = response_data["task_instances"][0]["id"]
_check_task_instance_note(session, ti_id, {"content":
"placeholder-note", "user_id": None})
+ def _seed_task_state(self, session, dag_id):
+ """Store one task state key for the single TI created by these
tests."""
+ ti = session.scalars(select(TaskInstance).where(TaskInstance.dag_id ==
dag_id)).one()
+ MetastoreBackend().set(
+ TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id,
map_index=ti.map_index),
+ "job_id",
+ "app_1234",
+ session=session,
+ )
+ session.commit()
+
+ def _task_state_rows(self, session, dag_id):
+ return
session.scalars(select(TaskStateStoreModel).where(TaskStateStoreModel.dag_id ==
dag_id)).all()
+
+ @pytest.mark.db_test
+ @pytest.mark.parametrize(
+ ("payload_extra", "expect_kept"),
+ [
+ pytest.param({}, False, id="default-discards"),
+ pytest.param({"keep_task_state": True}, True, id="keep-preserves"),
+ pytest.param({"keep_task_state": False}, False,
id="explicit-false-discards"),
Review Comment:
Fixed. `test_clear_task_state_store` now seeds two mapped indices of times_2
(0 and 1), clears only index 0, and asserts index 1's state row survives
untouched, on top of the discard/keep assertion for the targeted index.
_seed_task_state filters by task_id/map_index so .one() stays correct with
multiple mapped TIs.
--
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]
