amoghrajesh commented on code in PR #69874:
URL: https://github.com/apache/airflow/pull/69874#discussion_r3603184256
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py:
##########
@@ -1594,20 +1594,139 @@ def test_patch_dag_run_bad_request(self, test_client):
assert body["detail"][0]["msg"] == "Input should be 'queued',
'success' or 'failed'"
@pytest.mark.parametrize(
- ("state", "listener_state"),
+ ("state", "expected_dagrun_state"),
[
- ("queued", []),
- ("success", [DagRunState.SUCCESS]),
- ("failed", [DagRunState.FAILED]),
+ ("queued", None),
+ ("success", DagRunState.SUCCESS),
+ ("failed", DagRunState.FAILED),
],
)
@pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
- def test_patch_dag_run_notifies_listeners(self, test_client, state,
listener_state, listener_manager):
+ def test_patch_dag_run_notifies_listeners(
+ self, test_client, state, expected_dagrun_state, listener_manager
+ ):
listener = ClassBasedListener()
listener_manager(listener)
response =
test_client.patch(f"/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}", json={"state":
state})
assert response.status_code == 200
- assert listener.state == listener_state
+ assert listener.state == ([expected_dagrun_state] if
expected_dagrun_state else [])
+
+ @pytest.mark.parametrize(
+ ("dag_run_state", "expected_ti_state"),
+ [
+ ("success", TaskInstanceState.SUCCESS),
+ ("failed", TaskInstanceState.FAILED),
+ ],
+ )
+ @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
+ def test_patch_dag_run_notifies_ti_listeners_for_running_tasks(
+ self,
+ test_client,
+ dag_maker,
+ session,
+ listener_manager,
+ dag_run_state,
+ expected_ti_state,
+ ):
+ with dag_maker(dag_id="test_ti_listeners", schedule=None,
serialized=True):
+ EmptyOperator(task_id="t1")
+
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == "t1",
+ )
+ )
+ ti.state = TaskInstanceState.RUNNING
+ dag_maker.sync_dagbag_to_db()
+ session.commit()
+
+ listener = ClassBasedListener()
+ listener_manager(listener)
+
+ response = test_client.patch(
+ f"/dags/test_ti_listeners/dagRuns/{dr.run_id}", json={"state":
dag_run_state}
+ )
+ assert response.status_code == 200
+ # Use `is` for type-safe comparison between TaskInstanceState and
DagRunState.
Review Comment:
```suggestion
```
Announcing the assertion, not needed.
--
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]
