@@ -25,6 +25,22 @@ def assert_metric(role, role_group, metric):
2525 return metric in metric_response .text
2626
2727
28+ # Check if dag run state is "success", "queued", or "running"
29+ # TODO: in future, we could wait on it.
30+ # See: https://airflow.apache.org/docs/apache-airflow/3.1.6/stable-rest-api-ref.html#operation/wait_dag_run_until_finished
31+ def assert_dag_started (rest_url , headers , dag_id , dag_run_id ):
32+ dag_run_response = requests .get (
33+ f"{ rest_url } /dags/{ dag_id } /dagRuns/{ dag_run_id } " , headers = headers
34+ )
35+ dag_run_state = dag_run_response .json ()["state" ]
36+ print (f"DAG RUN STATE: { dag_run_state } " )
37+ return (
38+ dag_run_state == "success"
39+ or dag_run_state == "queued"
40+ or dag_run_state == "running"
41+ )
42+
43+
2844def metrics_v3 (role_group : str ) -> None :
2945 now = datetime .now (timezone .utc )
3046 ts = now .strftime ("%Y-%m-%dT%H:%M:%S.%f" ) + now .strftime ("%z" )
@@ -64,11 +80,16 @@ def metrics_v3(role_group: str) -> None:
6480 response = requests .patch (
6581 f"{ rest_url } /dags/{ dag_id } " , headers = headers , json = {"is_paused" : False }
6682 )
83+
6784 # trigger DAG
6885 response = requests .post (
6986 f"{ rest_url } /dags/{ dag_id } /dagRuns" , headers = headers , json = dag_data
7087 )
7188
89+ dag_run_id = response .json ()["dag_run_id" ]
90+
91+ print (f"DAG RUN ID: { dag_run_id } " )
92+
7293 # Test the DAG in a loop. Each time we call the script a new job will be started: we can avoid
7394 # or minimize this by looping over the check instead.
7495 iterations = 4
@@ -77,24 +98,13 @@ def metrics_v3(role_group: str) -> None:
7798 assert response .status_code == 200 , "DAG run could not be triggered."
7899 # Wait for the metrics to be consumed by the statsd-exporter
79100 time .sleep (5 )
80- # (disable line-break flake checks)
101+ heartbeat_metric = "airflow_scheduler_heartbeat"
102+ dag_run_success_count_metric = f"airflow_dagrun_duration_success_{ dag_id } _count"
81103 if (
82- (assert_metric ("scheduler" , role_group , "airflow_scheduler_heartbeat" ))
83- and (
84- assert_metric (
85- "webserver" ,
86- role_group ,
87- "airflow_task_instance_created_CoreDeferrableSleepOperator" ,
88- )
89- ) # noqa: W503, W504
90- and (
91- assert_metric (
92- "scheduler" ,
93- role_group ,
94- "airflow_dagrun_duration_success_core_deferrable_sleep_demo_count" ,
95- )
96- )
97- ): # noqa: W503, W504
104+ assert_dag_started (rest_url , headers , dag_id , dag_run_id )
105+ and assert_metric ("scheduler" , role_group , heartbeat_metric )
106+ and assert_metric ("scheduler" , role_group , dag_run_success_count_metric )
107+ ):
98108 break
99109 time .sleep (10 )
100110 loop += 1
0 commit comments