From b1aa5b28ab83c99beb9302c4b14321285dbc4303 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Fri, 21 Jul 2023 15:04:22 +0100 Subject: [PATCH] Fix broken test --- tests/airflow/test_graph.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/airflow/test_graph.py b/tests/airflow/test_graph.py index 803c0e018..8e8266d82 100644 --- a/tests/airflow/test_graph.py +++ b/tests/airflow/test_graph.py @@ -13,7 +13,7 @@ create_test_task_metadata, calculate_operator_class, ) -from cosmos.constants import ExecutionMode, DbtResourceType +from cosmos.constants import ExecutionMode, DbtResourceType, TestBehavior from cosmos.dbt.graph import DbtNode @@ -75,7 +75,7 @@ def test_build_airflow_graph_with_after_each(): dag=dag, execution_mode=ExecutionMode.LOCAL, task_args=task_args, - test_behavior="after_each", + test_behavior=TestBehavior.AFTER_EACH, dbt_project_name="astro_shop", conn_id="fake_conn", ) @@ -117,7 +117,7 @@ def test_build_airflow_graph_with_after_all(): dag=dag, execution_mode=ExecutionMode.LOCAL, task_args=task_args, - test_behavior="after_all", + test_behavior=TestBehavior.AFTER_ALL, dbt_project_name="astro_shop", conn_id="fake_conn", )