This is an automated email from the ASF dual-hosted git repository. ash pushed a commit to branch task-sdk-first-code in repository https://gitbox.apache.org/repos/asf/airflow.git
commit acfbf05e4750f65f4906cc9bdb29ea2039db08ff Author: Kaxil Naik <[email protected]> AuthorDate: Tue Oct 29 13:51:02 2024 +0000 Fix AirflowExecption error in tests/models/test_taskinstance.py --- providers/tests/standard/operators/test_weekday.py | 3 +-- task_sdk/src/airflow/sdk/definitions/node.py | 3 +-- tests/models/test_taskinstance.py | 10 +++++----- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/providers/tests/standard/operators/test_weekday.py b/providers/tests/standard/operators/test_weekday.py index 87ceed00ef6..c9df8b0124f 100644 --- a/providers/tests/standard/operators/test_weekday.py +++ b/providers/tests/standard/operators/test_weekday.py @@ -22,7 +22,6 @@ import datetime import pytest import time_machine -from airflow.exceptions import AirflowException from airflow.models.dagrun import DagRun from airflow.models.taskinstance import TaskInstance as TI from airflow.models.xcom import XCom @@ -208,7 +207,7 @@ class TestBranchDayOfWeekOperator: def test_branch_with_no_weekday(self, dag_maker): """Check if BranchDayOfWeekOperator raises exception on missing weekday""" - with pytest.raises(AirflowException): + with pytest.raises(TypeError, match="missing keyword argument 'week_day'"): with dag_maker( "branch_day_of_week_operator_test", start_date=DEFAULT_DATE, diff --git a/task_sdk/src/airflow/sdk/definitions/node.py b/task_sdk/src/airflow/sdk/definitions/node.py index b3b519a07f0..a29a36a7b7a 100644 --- a/task_sdk/src/airflow/sdk/definitions/node.py +++ b/task_sdk/src/airflow/sdk/definitions/node.py @@ -27,7 +27,6 @@ from typing import TYPE_CHECKING, Any import methodtools import re2 -from airflow.exceptions import AirflowException from airflow.sdk.definitions.mixins import DependencyMixin if TYPE_CHECKING: @@ -146,7 +145,7 @@ class DAGNode(DependencyMixin, metaclass=ABCMeta): dags: set[DAG] = {task.dag for task in [*self.roots, *task_list] if task.has_dag() and task.dag} if len(dags) > 1: - raise AirflowException(f"Tried to set relationships between tasks in more than one DAG: {dags}") + raise RuntimeError(f"Tried to set relationships between tasks in more than one DAG: {dags}") elif len(dags) == 1: dag = dags.pop() else: diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index de90889cffd..ef213f43837 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -229,7 +229,7 @@ class TestTaskInstance: # no dag assigned assert not op.has_dag() - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): getattr(op, "dag") # no improper assignment @@ -239,7 +239,7 @@ class TestTaskInstance: op.dag = dag # no reassignment - with pytest.raises(AirflowException): + with pytest.raises(ValueError): op.dag = dag2 # but assigning the same dag is ok @@ -261,7 +261,7 @@ class TestTaskInstance: assert [i.has_dag() for i in [op1, op2, op3, op4]] == [False, False, True, True] # can't combine operators with no dags - with pytest.raises(AirflowException): + with pytest.raises(ValueError): op1.set_downstream(op2) # op2 should infer dag from op1 @@ -270,9 +270,9 @@ class TestTaskInstance: assert op2.dag is dag # can't assign across multiple DAGs - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): op1.set_downstream(op4) - with pytest.raises(AirflowException): + with pytest.raises(RuntimeError): op1.set_downstream([op3, op4]) def test_bitshift_compose_operators(self, dag_maker):
