Repository: astronomer/astronomer-cosmos
Pull request #1743: Run listener tests in Airflow 3
Author: login: pankajastro
Labels: name: area:testing
description: Related to testing, like unit tests, integration tests, etc
name: size:M
description: This PR changes 30-99 lines, ignoring generated files.
name: lgtm
description: This PR has been approved by a maintainer
Created: 2025-05-05T15:32:17Z | Merged: 2025-05-06T10:37:49Z

closes: https://github.com/astronomer/astronomer-cosmos/issues/1703

Previously, we sent the dag_hash to Scarf,
but with the removal of DagRun.dag_hash in Airflow 3,
this PR modifies the implementation to send a portion 
of the dag_id hash instead. This change is a reasonable 
compromise, as we only need a unique identifier for the DAG. 
Additionally, the PR updates the relevant test cases to reflect this change.




Diff:
diff --git a/cosmos/listeners/dag_run_listener.py b/cosmos/listeners/dag_run_listener.py
index 0314c3474d..2abd21e7d1 100644
--- a/cosmos/listeners/dag_run_listener.py
+++ b/cosmos/listeners/dag_run_listener.py
@@ -1,12 +1,19 @@
 from __future__ import annotations
 
+import hashlib
+
+from airflow import __version__ as airflow_version
 from airflow.listeners import hookimpl
 from airflow.models.dag import DAG
 from airflow.models.dagrun import DagRun
+from packaging import version
 
 from cosmos import telemetry
+from cosmos.constants import _AIRFLOW3_MAJOR_VERSION
 from cosmos.log import get_logger
 
+AIRFLOW_VERSION_MAJOR = version.parse(airflow_version).major
+
 logger = get_logger(__name__)
 
 
@@ -50,8 +57,13 @@ def on_dag_run_success(dag_run: DagRun, msg: str) -> None:
         logger.debug("The DAG does not use Cosmos")
         return
 
+    if AIRFLOW_VERSION_MAJOR < _AIRFLOW3_MAJOR_VERSION:
+        dag_hash = dag_run.dag_hash
+    else:
+        dag_hash = hashlib.md5(dag_run.dag_id.encode("utf-8")).hexdigest()
+
     additional_telemetry_metrics = {
-        "dag_hash": dag_run.dag_hash,
+        "dag_hash": dag_hash,
         "status": EventStatus.SUCCESS,
         "task_count": len(serialized_dag.task_ids),
         "cosmos_task_count": total_cosmos_tasks(serialized_dag),
@@ -73,8 +85,13 @@ def on_dag_run_failed(dag_run: DagRun, msg: str) -> None:
         logger.debug("The DAG does not use Cosmos")
         return
 
+    if AIRFLOW_VERSION_MAJOR < _AIRFLOW3_MAJOR_VERSION:
+        dag_hash = dag_run.dag_hash
+    else:
+        dag_hash = hashlib.md5(dag_run.dag_id.encode("utf-8")).hexdigest()
+
     additional_telemetry_metrics = {
-        "dag_hash": dag_run.dag_hash,
+        "dag_hash": dag_hash,
         "status": EventStatus.FAILED,
         "task_count": len(serialized_dag.task_ids),
         "cosmos_task_count": total_cosmos_tasks(serialized_dag),
diff --git a/tests/listeners/test_dag_run_listener.py b/tests/listeners/test_dag_run_listener.py
index 77e30f8466..456dea9c1a 100644
--- a/tests/listeners/test_dag_run_listener.py
+++ b/tests/listeners/test_dag_run_listener.py
@@ -1,6 +1,6 @@
 import logging
 import uuid
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from unittest.mock import patch
 
@@ -13,12 +13,15 @@
 from cosmos import DbtRunLocalOperator, ProfileConfig, ProjectConfig
 from cosmos.airflow.dag import DbtDag
 from cosmos.airflow.task_group import DbtTaskGroup
+from cosmos.constants import _AIRFLOW3_MAJOR_VERSION
 from cosmos.listeners.dag_run_listener import on_dag_run_failed, on_dag_run_success, total_cosmos_tasks
 from cosmos.profiles import PostgresUserPasswordProfileMapping
 
 DBT_ROOT_PATH = Path(__file__).parent.parent.parent / "dev/dags/dbt"
 DBT_PROJECT_NAME = "jaffle_shop"
 
+AIRFLOW_VERSION_MAJOR = version.parse(airflow_version).major
+
 profile_config = ProfileConfig(
     profile_name="default",
     target_name="dev",
@@ -79,8 +82,6 @@ def test_not_cosmos_dag():
     assert total_cosmos_tasks(dag) == 0
 
 
-# TODO: Make test compatible with Airflow 3.0. Issue:https://github.com/astronomer/astronomer-cosmos/issues/1703
-@pytest.mark.skipif(version.parse(airflow_version).major == 3, reason="Test need to be updated for Airflow 3.0")
 @pytest.mark.integration
 @patch("cosmos.listeners.dag_run_listener.telemetry.emit_usage_metrics_if_enabled")
 def test_on_dag_run_success(mock_emit_usage_metrics_if_enabled, caplog):
@@ -95,10 +96,25 @@ def test_on_dag_run_success(mock_emit_usage_metrics_if_enabled, caplog):
         dag_id="basic_cosmos_dag",
     )
     run_id = str(uuid.uuid1())
-    dag_run = dag.create_dagrun(
-        state=State.NONE,
-        run_id=run_id,
-    )
+
+    run_after = datetime.now(timezone.utc) - timedelta(seconds=1)
+    if AIRFLOW_VERSION_MAJOR < _AIRFLOW3_MAJOR_VERSION:
+        # Airflow 2
+        dag_run = dag.create_dagrun(
+            state=State.NONE,
+            run_id=run_id,
+        )
+    else:
+        # Airflow 3
+        from airflow.utils.types import DagRunTriggeredByType, DagRunType
+
+        dag_run = dag.create_dagrun(
+            state=State.NONE,
+            run_id=run_id,
+            run_after=run_after,
+            run_type=DagRunType.MANUAL,
+            triggered_by=DagRunTriggeredByType.TIMETABLE,
+        )
 
     on_dag_run_success(dag_run, msg="test success")
     assert "Running on_dag_run_success" in caplog.text
@@ -106,8 +122,6 @@ def test_on_dag_run_success(mock_emit_usage_metrics_if_enabled, caplog):
     assert mock_emit_usage_metrics_if_enabled.call_count == 1
 
 
-# TODO: Make test compatible with Airflow 3.0. Issue:https://github.com/astronomer/astronomer-cosmos/issues/1703
-@pytest.mark.skipif(version.parse(airflow_version).major == 3, reason="Test need to be updated for Airflow 3.0")
 @pytest.mark.integration
 @patch("cosmos.listeners.dag_run_listener.telemetry.emit_usage_metrics_if_enabled")
 def test_on_dag_run_failed(mock_emit_usage_metrics_if_enabled, caplog):
@@ -122,10 +136,24 @@ def test_on_dag_run_failed(mock_emit_usage_metrics_if_enabled, caplog):
         dag_id="basic_cosmos_dag",
     )
     run_id = str(uuid.uuid1())
-    dag_run = dag.create_dagrun(
-        state=State.FAILED,
-        run_id=run_id,
-    )
+    run_after = datetime.now(timezone.utc) - timedelta(seconds=1)
+    if AIRFLOW_VERSION_MAJOR < _AIRFLOW3_MAJOR_VERSION:
+        # Airflow 2
+        dag_run = dag.create_dagrun(
+            state=State.NONE,
+            run_id=run_id,
+        )
+    else:
+        # Airflow 3
+        from airflow.utils.types import DagRunTriggeredByType, DagRunType
+
+        dag_run = dag.create_dagrun(
+            state=State.NONE,
+            run_id=run_id,
+            run_after=run_after,
+            run_type=DagRunType.MANUAL,
+            triggered_by=DagRunTriggeredByType.TIMETABLE,
+        )
 
     on_dag_run_failed(dag_run, msg="test failed")
     assert "Running on_dag_run_failed" in caplog.text


Reviews:
user: login: pankajastro
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T15:33:15Z
user: login: pankajastro
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T15:34:25Z
user: login: pankajastro
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T15:35:02Z
user: login: pankajastro
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T15:35:27Z
user: login: tatiana
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T08:35:57Z
user: login: pankajastro
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-06T09:55:02Z
user: login: tatiana
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T10:36:41Z
user: login: tatiana
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-06T10:36:54Z

Review comments:
diff_hunk: @@ -2,7 +2,7 @@ name: test
 
 on:
   push: # Run on pushes to the default branch
-    branches: [main]
+    branches: [main, fix_af3_test]
path: .github/workflows/test.yml
user: login: pankajastro
body: ```suggestion
    branches: [main]
```
created_at: 2025-05-05T15:33:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -79,8 +82,8 @@ def test_not_cosmos_dag():
     assert total_cosmos_tasks(dag) == 0
 
 
-# TODO: Make test compatible with Airflow 3.0. Issue:https://github.com/astronomer/astronomer-cosmos/issues/1703
-@pytest.mark.skipif(version.parse(airflow_version).major == 3, reason="Test need to be updated for Airflow 3.0")
+# # TODO: Make test compatible with Airflow 3.0. Issue:https://github.com/astronomer/astronomer-cosmos/issues/1703
+# @pytest.mark.skipif(version.parse(airflow_version).major == 3, reason="Test need to be updated for Airflow 3.0")
path: tests/listeners/test_dag_run_listener.py
user: login: pankajastro
body: ```suggestion
```
created_at: 2025-05-05T15:34:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -107,7 +123,7 @@ def test_on_dag_run_success(mock_emit_usage_metrics_if_enabled, caplog):
 
 
 # TODO: Make test compatible with Airflow 3.0. Issue:https://github.com/astronomer/astronomer-cosmos/issues/1703
-@pytest.mark.skipif(version.parse(airflow_version).major == 3, reason="Test need to be updated for Airflow 3.0")
+# @pytest.mark.skipif(version.parse(airflow_version).major == 3, reason="Test need to be updated for Airflow 3.0")
path: tests/listeners/test_dag_run_listener.py
user: login: pankajastro
body: ```suggestion
```
created_at: 2025-05-05T15:35:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -107,7 +123,6 @@ def test_on_dag_run_success(mock_emit_usage_metrics_if_enabled, caplog):
 
 
 # TODO: Make test compatible with Airflow 3.0. Issue:https://github.com/astronomer/astronomer-cosmos/issues/1703
path: tests/listeners/test_dag_run_listener.py
user: login: pankajastro
body: ```suggestion
```
created_at: 2025-05-05T15:35:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -51,7 +53,7 @@ def on_dag_run_success(dag_run: DagRun, msg: str) -> None:
         return
 
     additional_telemetry_metrics = {
-        "dag_hash": dag_run.dag_hash,
+        "dag_hash": hashlib.md5(dag_run.dag_id.encode("utf-8")).hexdigest()[:8],
path: cosmos/listeners/dag_run_listener.py
user: login: tatiana
body: @pankajastro, do we know how this hash is currently generated? How big is Airflow's original hash, and how likely are we to have collisions?

Also, should we only change this for Airflow 3, so we don't mess up the existing metrics in Scarf for AF2 users?
created_at: 2025-05-06T08:35:57Z
author_association: COLLABORATOR
diff_hunk: @@ -51,7 +53,7 @@ def on_dag_run_success(dag_run: DagRun, msg: str) -> None:
         return
 
     additional_telemetry_metrics = {
-        "dag_hash": dag_run.dag_hash,
+        "dag_hash": hashlib.md5(dag_run.dag_id.encode("utf-8")).hexdigest()[:8],
path: cosmos/listeners/dag_run_listener.py
user: login: pankajastro
body: The dag_hash param has been removed, and I can't see an equivalent param in Airflow 3. The possibility of collision is rare in practice, and Airflow uses MD5 for hashing.

I have restored the previous behaviour for AF2.

created_at: 2025-05-06T09:55:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -51,7 +53,7 @@ def on_dag_run_success(dag_run: DagRun, msg: str) -> None:
         return
 
     additional_telemetry_metrics = {
-        "dag_hash": dag_run.dag_hash,
+        "dag_hash": hashlib.md5(dag_run.dag_id.encode("utf-8")).hexdigest()[:8],
path: cosmos/listeners/dag_run_listener.py
user: login: tatiana
body: Thanks, @pankajastro , looks good!
created_at: 2025-05-06T10:36:41Z
author_association: COLLABORATOR

PR comments:
user: login: codecov[bot]
created_at: 2025-05-05T15:32:35Z
body: ## [Codecov](https://app.codecov.io/gh/astronomer/astronomer-cosmos/pull/1743?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astronomer) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 97.70%. Comparing base [(`e3b0b5e`)](https://app.codecov.io/gh/astronomer/astronomer-cosmos/commit/e3b0b5e213a174e6dac107814a75353b53c6bd8d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astronomer) to head [(`650935f`)](https://app.codecov.io/gh/astronomer/astronomer-cosmos/commit/650935f1c85f765396713a6b4314c826b75014af?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astronomer).
> Report is 2 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #1743   +/-   ##
=======================================
  Coverage   97.69%   97.70%           
=======================================
  Files          84       84           
  Lines        5209     5220   +11     
=======================================
+ Hits         5089     5100   +11     
  Misses        120      120           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/astronomer/astronomer-cosmos/pull/1743?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astronomer).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=astronomer).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: netlify[bot]
created_at: 2025-05-05T15:32:46Z
body: ### <span aria-hidden="true">✅</span> Deploy Preview for *sunny-pastelito-5ecb04* canceled.


|  Name | Link |
|:-:|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | 650935f1c85f765396713a6b4314c826b75014af |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/6819e642459a80000848eeff |
author_association: NONE
user: login: cloudflare-workers-and-pages[bot]
created_at: 2025-05-05T15:33:25Z
body: ## Deploying astronomer-cosmos with &nbsp;<a href="https://pages.dev"><img alt="Cloudflare Pages" src="https://user-images.githubusercontent.com/23264/106598434-9e719e00-654f-11eb-9e59-6167043cfa01.png" width="16"></a> &nbsp;Cloudflare Pages

<table><tr><td><strong>Latest commit:</strong> </td><td>
<code>650935f</code>
</td></tr>
<tr><td><strong>Status:</strong></td><td>&nbsp;✅&nbsp; Deploy successful!</td></tr>
<tr><td><strong>Preview URL:</strong></td><td>
<a href='https://2efac587.astronomer-cosmos.pages.dev'>https://2efac587.astronomer-cosmos.pages.dev</a>
</td></tr>
<tr><td><strong>Branch Preview URL:</strong></td><td>
<a href='https://fix-af3-test.astronomer-cosmos.pages.dev'>https://fix-af3-test.astronomer-cosmos.pages.dev</a>
</td></tr>
</table>

[View logs](https://dash.cloudflare.com/799f6c6de390adc4994d0b9621414a76/pages/view/astronomer-cosmos/2efac587-54f9-4ada-ae63-3bb61a0bc86b)

author_association: NONE