verify test parameters

Author: Bowrna

airflow-core/tests/unit/decorators/test_signature_consistency.py

@@ -0,0 +1,106 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import inspect
+import os
+
+import libcst as cst
+import pytest
+
+from airflow.decorators import task
+
+DECORATOR_OPERATOR_MAP = {
+    "kubernetes": "airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator",
+    "sensor": "airflow.sdk.bases.sensor.BaseSensorOperator",
+    "virtualenv": "airflow.providers.standard.operators.python.PythonVirtualenvOperator",
+    "branch_virtualenv": "airflow.providers.standard.operators.python.BranchPythonVirtualenvOperator",
+    # Add more here...
+}
+
+
+def extract_function_params(code, function_name, return_type):
+    """Extracts parameters from a specific function definition in the given code.
+
+    Args:
+        code (str): The Python code to parse.
+        function_name (str): The name of the function to extract parameters from.
+        return_type (str): As the pyi file has multiple @overload decorator, extract function param based on return type.
+
+    Returns:
+        list: A list of parameter names, or None if the function is not found.
+    """
+    module = cst.parse_module(code)
+
+    class FunctionParamExtractor(cst.CSTVisitor):
+        def __init__(self, target_function_name, target_return_type):
+            self.target_function_name = target_function_name
+            self.target_return_type = target_return_type
+            self.params: list[str] = []
+
+        def visit_FunctionDef(self, node):
+            # Match function name
+            if node.name.value == self.target_function_name:
+                if node.returns:
+                    annotation = node.returns.annotation
+                    if isinstance(annotation, cst.Name) and annotation.value == self.target_return_type:
+                        parameters_node = node.params
+                        self.params.extend(param.name.value for param in parameters_node.params)
+                        self.params.extend(param.name.value for param in parameters_node.kwonly_params)
+                        self.params.extend(param.name.value for param in parameters_node.posonly_params)
+                        if parameters_node.star_kwarg:
+                            self.params.append(parameters_node.star_kwarg.name.value)
+                            return False  # Stop traversing after finding the real function
+            return True  # Keep traversing
+
+    extractor = FunctionParamExtractor(function_name, return_type)
+    module.visit(extractor)
+    return extractor.params
+
+
+def get_decorator_params(decorator_name: str):
+    file_path = os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "../../../../task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi",
+        )
+    )
+    params = []
+    with open(file_path) as file:
+        code = file.read()
+        params = extract_function_params(code, decorator_name, "TaskDecorator")
+    return set(params)
+
+
+def get_operator_params(operator_path: str):
+    module_path, class_name = operator_path.rsplit(".", 1)
+    module = __import__(module_path, fromlist=[class_name])
+    operator_cls = getattr(module, class_name)
+    sig = inspect.signature(operator_cls.__init__)
+    return set(p for p in sig.parameters.keys() if p not in ("self", "args", "kwargs"))
+
+
+@pytest.mark.parametrize("decorator, operator_path", DECORATOR_OPERATOR_MAP.items())
+def test_decorator_matches_operator_signature(decorator, operator_path):
+    decorator_params = get_decorator_params(decorator)
+    operator_params = get_operator_params(operator_path)
+    missing_in_decorator = operator_params - decorator_params
+
+    ignored = {"kwargs", "args", "self"}
+    missing_in_decorator -= ignored
+    assert not missing_in_decorator, f"{decorator} is missing params: {missing_in_decorator}"

task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi

@@ -27,6 +27,7 @@ from typing import Any, Callable, TypeVar, overload
 from docker.types import Mount
 from kubernetes.client import models as k8s
 
+from airflow.providers.cncf.kubernetes.callbacks import KubernetesPodOperatorCallback
 from airflow.providers.cncf.kubernetes.secret import Secret
 from airflow.sdk.bases.decorator import FParams, FReturn, Task, TaskDecorator, _TaskDecorator
 from airflow.sdk.definitions.dag import dag
@@ -96,16 +97,21 @@ class TaskDecoratorCollection:
         multiple_outputs: bool | None = None,
         # 'python_callable', 'op_args' and 'op_kwargs' since they are filled by
         # _PythonVirtualenvDecoratedOperator.
+        python_callable: Callable,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
         requirements: None | Iterable[str] | str = None,
         python_version: None | str | int | float = None,
         serializer: Literal["pickle", "cloudpickle", "dill"] | None = None,
         system_site_packages: bool = True,
         templates_dict: Mapping[str, Any] | None = None,
+        templates_exts: list[str] | None = None,
         pip_install_options: list[str] | None = None,
+        expect_airflow: bool = True,
         skip_on_exit_code: int | Container[int] | None = None,
         index_urls: None | Collection[str] | str = None,
         venv_cache_path: None | str = None,
-        show_return_value_in_logs: bool = True,
+        string_args: Iterable[str] | None = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
         **kwargs,
@@ -218,18 +224,24 @@ class TaskDecoratorCollection:
         self,
         *,
         multiple_outputs: bool | None = None,
-        # 'python_callable', 'op_args' and 'op_kwargs' since they are filled by
-        # _PythonVirtualenvDecoratedOperator.
+        python_callable: Callable,
         requirements: None | Iterable[str] | str = None,
+        op_args: Collection[Any] | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        string_args: Iterable[str] | None = None,
         python_version: None | str | int | float = None,
         serializer: Literal["pickle", "cloudpickle", "dill"] | None = None,
         system_site_packages: bool = True,
         templates_dict: Mapping[str, Any] | None = None,
+        templates_exts: list[str] | None = None,
         pip_install_options: list[str] | None = None,
         skip_on_exit_code: int | Container[int] | None = None,
         index_urls: None | Collection[str] | str = None,
         venv_cache_path: None | str = None,
+        expect_airflow: bool = True,
         show_return_value_in_logs: bool = True,
+        env_vars: dict[str, str] | None = None,
+        inherit_env: bool = True,
         **kwargs,
     ) -> TaskDecorator:
         """Create a decorator to wrap the decorated callable into a BranchPythonVirtualenvOperator.
@@ -507,6 +519,7 @@ class TaskDecoratorCollection:
         image: str | None = None,
         name: str | None = None,
         random_name_suffix: bool = ...,
+        cmds: list[str] | None = None,
         arguments: list[str] | None = None,
         ports: list[k8s.V1ContainerPort] | None = None,
         volume_mounts: list[k8s.V1VolumeMount] | None = None,
@@ -520,6 +533,7 @@ class TaskDecoratorCollection:
         reattach_on_restart: bool = ...,
         startup_timeout_seconds: int = ...,
         startup_check_interval_seconds: int = ...,
+        schedule_timeout_seconds: int | None = None,
         get_logs: bool = True,
         container_logs: Iterable[str] | str | Literal[True] = ...,
         image_pull_policy: str | None = None,
@@ -530,6 +544,7 @@ class TaskDecoratorCollection:
         node_selector: dict | None = None,
         image_pull_secrets: list[k8s.V1LocalObjectReference] | None = None,
         service_account_name: str | None = None,
+        automount_service_account_token: bool | None = None,
         hostnetwork: bool = False,
         host_aliases: list[k8s.V1HostAlias] | None = None,
         tolerations: list[k8s.V1Toleration] | None = None,
@@ -553,13 +568,19 @@ class TaskDecoratorCollection:
         skip_on_exit_code: int | Container[int] | None = None,
         base_container_name: str | None = None,
         base_container_status_polling_interval: float = ...,
+        init_container_logs: Iterable[str] | str | Literal[True] | None = None,
         deferrable: bool = ...,
         poll_interval: float = ...,
         log_pod_spec_on_failure: bool = ...,
         on_finish_action: str = ...,
+        is_delete_operator_pod: None | bool = None,
         termination_message_policy: str = ...,
         active_deadline_seconds: int | None = None,
+        callbacks: (
+            list[type[KubernetesPodOperatorCallback]] | type[KubernetesPodOperatorCallback] | None
+        ) = None,
         progress_callback: Callable[[str], None] | None = None,
+        logging_interval: int | None = None,
         **kwargs,
     ) -> TaskDecorator:
         """Create a decorator to convert a callable to a Kubernetes Pod task.
@@ -849,6 +870,8 @@ class TaskDecoratorCollection:
         mode: str = ...,
         exponential_backoff: bool = False,
         max_wait: timedelta | float | None = None,
+        silent_fail: bool = False,
+        never_fail: bool = False,
         **kwargs,
     ) -> TaskDecorator:
         """
@@ -873,6 +896,13 @@ class TaskDecoratorCollection:
         :param exponential_backoff: allow progressive longer waits between
             pokes by using exponential backoff algorithm
         :param max_wait: maximum wait interval between pokes, can be ``timedelta`` or ``float`` seconds
+        :param silent_fail: If true, and poke method raises an exception different from
+            AirflowSensorTimeout, AirflowTaskTimeout, AirflowSkipException
+            and AirflowFailException, the sensor will log the error and continue
+            its execution. Otherwise, the sensor task fails, and it can be retried
+            based on the provided `retries` parameter.
+        :param never_fail: If true, and poke method raises an exception, sensor will be skipped.
+           Mutually exclusive with soft_fail.
         """
     @overload
     def sensor(self, python_callable: Callable[FParams, FReturn] | None = None) -> Task[FParams, FReturn]: ...