From 98d619d3c5fb01effee3c68e15d45b53c2e9959a Mon Sep 17 00:00:00 2001 From: phi Date: Fri, 6 Sep 2024 22:13:42 +0900 Subject: [PATCH] fix: rm check get_python_source --- tests/utils/test_decorators.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/utils/test_decorators.py b/tests/utils/test_decorators.py index 82c04bcd535b..19d3ec31d031 100644 --- a/tests/utils/test_decorators.py +++ b/tests/utils/test_decorators.py @@ -24,7 +24,12 @@ if TYPE_CHECKING: from airflow.decorators.base import Task, TaskDecorator -DECORATORS = sorted(set(x for x in dir(task) if not x.startswith("_")) - {"skip_if", "run_if"}) + +_CONDITION_DECORATORS = frozenset({"skip_if", "run_if"}) +_NO_SOURCE_DECORATORS = frozenset({"sensor"}) +DECORATORS = sorted( + set(x for x in dir(task) if not x.startswith("_")) - _CONDITION_DECORATORS - _NO_SOURCE_DECORATORS +) DECORATORS_USING_SOURCE = ("external_python", "virtualenv", "branch_virtualenv", "branch_external_python") @@ -115,9 +120,6 @@ def f(): def parse_python_source(task: Task, custom_operator_name: str | None = None) -> str: operator = task().operator - if not hasattr(operator, "get_python_source"): - pytest.skip(f"Operator {operator} does not have get_python_source method") - if custom_operator_name: custom_operator_name = ( custom_operator_name if custom_operator_name.startswith("@") else f"@{custom_operator_name}"