From 9c5d93b40cb092ae274a492866a0fd727db60982 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Thu, 6 Mar 2025 14:57:53 +0100 Subject: [PATCH] Move task_sdk to a standalone task-sdk distribution This is the next stage of refactoring of airflow packages, after moving providers to standalone dstribution and separating devel-common as a common distribution. The `task_sdk` has been renamed to `task-sdk` - this way we will never import anything in task_sdk accidentally starting from content root. Some changes have been needed to make it works: * autouse fixture was added to pytest plugin to add `task-sdk/tests` to PYTHONPATH to make it root import * all tests were moved to `task_sdk` package inside the tests folder * all imports for tests are now `from task_sdk` * common tools for task_sdk has been moved to `devel-common/src/test_utils/task_sdk.py` in order to allow importing them before `task-sdk/tests` is added to pythonpath --- .dockerignore | 2 +- .github/boring-cyborg.yml | 2 +- .github/dependabot.yml | 2 +- .pre-commit-config.yaml | 16 +++---- Dockerfile | 2 +- Dockerfile.ci | 4 +- airflow/models/taskinstance.py | 2 +- airflow/utils/context.py | 2 +- contributing-docs/07_local_virtualenv.rst | 6 +-- .../contributors_quick_start_pycharm.rst | 8 ++-- .../commands/release_management_commands.py | 6 +-- .../commands/testing_commands.py | 2 +- .../src/airflow_breeze/global_constants.py | 2 +- .../utils/docker_command_utils.py | 2 +- .../src/airflow_breeze/utils/run_tests.py | 4 +- .../airflow_breeze/utils/selective_checks.py | 22 ++++----- .../tests/test_pytest_args_for_test_types.py | 2 +- dev/breeze/tests/test_selective_checks.py | 14 +++--- .../logging-monitoring/callbacks.rst | 2 +- .../openlineage/extractors/test_manager.py | 26 ++++++++++- pyproject.toml | 4 +- scripts/ci/docker-compose/local.yml | 4 +- .../providers-and-tests-sources.yml | 8 ++-- scripts/ci/docker-compose/remove-sources.yml | 4 +- scripts/ci/kubernetes/k8s_requirements.txt | 2 +- .../base_operator_partial_arguments.py | 4 +- scripts/ci/pre_commit/mypy_folder.py | 6 +-- scripts/ci/pre_commit/sync_init_decorator.py | 2 +- .../pre_commit/template_context_key_sync.py | 4 +- scripts/ci/testing/run_unit_tests.sh | 2 +- scripts/docker/entrypoint_ci.sh | 2 +- scripts/docker/install_airflow.sh | 2 +- .../install_airflow_and_providers.py | 6 +-- {task_sdk => task-sdk}/README.md | 0 {task_sdk => task-sdk}/dev/generate_models.py | 0 {task_sdk => task-sdk}/pyproject.toml | 39 ++++++++++++++++ .../src/airflow/__init__.py | 0 .../src/airflow/sdk/__init__.py | 0 .../src/airflow/sdk/api}/__init__.py | 0 .../src/airflow/sdk/api/client.py | 0 .../airflow/sdk/api/datamodels}/__init__.py | 0 .../airflow/sdk/api/datamodels/_generated.py | 0 .../airflow/sdk/api/datamodels/activities.py | 0 .../src/airflow/sdk/definitions}/__init__.py | 0 .../sdk/definitions/_internal}/__init__.py | 0 .../definitions/_internal/abstractoperator.py | 0 .../definitions/_internal/contextmanager.py | 0 .../_internal/dag_parsing_context.py | 0 .../sdk/definitions/_internal/decorators.py | 0 .../sdk/definitions/_internal/expandinput.py | 0 .../sdk/definitions/_internal/mixins.py | 0 .../airflow/sdk/definitions/_internal/node.py | 0 .../sdk/definitions/_internal/templater.py | 0 .../sdk/definitions/_internal/types.py | 0 .../airflow/sdk/definitions/asset/__init__.py | 0 .../sdk/definitions/asset/decorators.py | 0 .../airflow/sdk/definitions/asset/metadata.py | 0 .../airflow/sdk/definitions/baseoperator.py | 0 .../src/airflow/sdk/definitions/connection.py | 0 .../src/airflow/sdk/definitions/context.py | 0 .../src/airflow/sdk/definitions/dag.py | 0 .../src/airflow/sdk/definitions/edges.py | 0 .../src/airflow/sdk/definitions/macros.py | 0 .../airflow/sdk/definitions/mappedoperator.py | 0 .../src/airflow/sdk/definitions/param.py | 0 .../src/airflow/sdk/definitions/taskgroup.py | 0 .../src/airflow/sdk/definitions/template.py | 0 .../src/airflow/sdk/definitions/variable.py | 0 .../src/airflow/sdk/definitions/xcom_arg.py | 0 .../src/airflow/sdk/exceptions.py | 0 .../airflow/sdk/execution_time/__init__.py | 0 .../src/airflow/sdk/execution_time/comms.py | 0 .../src/airflow/sdk/execution_time/context.py | 0 .../sdk/execution_time/execute_workload.py | 0 .../sdk/execution_time/lazy_sequence.py | 0 .../sdk/execution_time/secrets_masker.py | 0 .../airflow/sdk/execution_time/supervisor.py | 0 .../airflow/sdk/execution_time/task_runner.py | 0 {task_sdk => task-sdk}/src/airflow/sdk/log.py | 0 .../src/airflow/sdk/py.typed | 0 .../src/airflow/sdk/types.py | 0 .../_internal => task-sdk/tests}/__init__.py | 0 {task_sdk => task-sdk}/tests/conftest.py | 2 +- task-sdk/tests/task_sdk/__init__.py | 45 +++++++++++++++++++ .../tests/task_sdk/api}/__init__.py | 0 .../tests/task_sdk}/api/test_client.py | 24 +--------- .../task_sdk}/dags/dag_parsing_context.py | 0 .../tests/task_sdk}/dags/super_basic.py | 0 .../dags/super_basic_deferred_run.py | 0 .../tests/task_sdk}/dags/super_basic_run.py | 0 .../tests/task_sdk/definitions}/__init__.py | 0 .../definitions/_internal}/__init__.py | 0 .../definitions/_internal/test_templater.py | 0 .../tests/task_sdk}/definitions/conftest.py | 0 .../tests/task_sdk}/definitions/test_asset.py | 0 .../definitions/test_asset_decorators.py | 0 .../definitions/test_baseoperator.py | 0 .../task_sdk}/definitions/test_connections.py | 0 .../task_sdk}/definitions/test_context.py | 0 .../tests/task_sdk}/definitions/test_dag.py | 0 .../task_sdk}/definitions/test_macros.py | 0 .../definitions/test_mappedoperator.py | 0 .../task_sdk}/definitions/test_mixins.py | 0 .../tests/task_sdk}/definitions/test_param.py | 0 .../definitions/test_secrets_masker.py | 0 .../task_sdk}/definitions/test_template.py | 0 .../task_sdk}/definitions/test_xcom_arg.py | 0 .../task_sdk/execution_time}/__init__.py | 0 .../task_sdk}/execution_time/conftest.py | 0 .../task_sdk}/execution_time/test_context.py | 0 .../execution_time/test_supervisor.py | 4 +- .../execution_time/test_task_runner.py | 3 +- .../tests/task_sdk/log}/test_log.py | 0 task_sdk/tests/execution_time/__init__.py | 16 ------- 114 files changed, 189 insertions(+), 122 deletions(-) rename {task_sdk => task-sdk}/README.md (100%) rename {task_sdk => task-sdk}/dev/generate_models.py (100%) rename {task_sdk => task-sdk}/pyproject.toml (80%) rename {task_sdk => task-sdk}/src/airflow/__init__.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/__init__.py (100%) rename {task_sdk => task-sdk/src/airflow/sdk/api}/__init__.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/api/client.py (100%) rename {task_sdk/src/airflow/sdk/api => task-sdk/src/airflow/sdk/api/datamodels}/__init__.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/api/datamodels/_generated.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/api/datamodels/activities.py (100%) rename {task_sdk/src/airflow/sdk/api/datamodels => task-sdk/src/airflow/sdk/definitions}/__init__.py (100%) rename {task_sdk/src/airflow/sdk/definitions => task-sdk/src/airflow/sdk/definitions/_internal}/__init__.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/abstractoperator.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/contextmanager.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/dag_parsing_context.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/decorators.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/expandinput.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/mixins.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/node.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/templater.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/_internal/types.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/asset/__init__.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/asset/decorators.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/asset/metadata.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/baseoperator.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/connection.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/context.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/dag.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/edges.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/macros.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/mappedoperator.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/param.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/taskgroup.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/template.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/variable.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/definitions/xcom_arg.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/exceptions.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/__init__.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/comms.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/context.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/execute_workload.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/lazy_sequence.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/secrets_masker.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/supervisor.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/execution_time/task_runner.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/log.py (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/py.typed (100%) rename {task_sdk => task-sdk}/src/airflow/sdk/types.py (100%) rename {task_sdk/src/airflow/sdk/definitions/_internal => task-sdk/tests}/__init__.py (100%) rename {task_sdk => task-sdk}/tests/conftest.py (99%) create mode 100644 task-sdk/tests/task_sdk/__init__.py rename {task_sdk/tests => task-sdk/tests/task_sdk/api}/__init__.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/api/test_client.py (97%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/dags/dag_parsing_context.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/dags/super_basic.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/dags/super_basic_deferred_run.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/dags/super_basic_run.py (100%) rename {task_sdk/tests/api => task-sdk/tests/task_sdk/definitions}/__init__.py (100%) rename {task_sdk/tests/definitions => task-sdk/tests/task_sdk/definitions/_internal}/__init__.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/_internal/test_templater.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/conftest.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_asset.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_asset_decorators.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_baseoperator.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_connections.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_context.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_dag.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_macros.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_mappedoperator.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_mixins.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_param.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_secrets_masker.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_template.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/definitions/test_xcom_arg.py (100%) rename {task_sdk/tests/definitions/_internal => task-sdk/tests/task_sdk/execution_time}/__init__.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/execution_time/conftest.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/execution_time/test_context.py (100%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/execution_time/test_supervisor.py (99%) rename {task_sdk/tests => task-sdk/tests/task_sdk}/execution_time/test_task_runner.py (99%) rename {task_sdk/tests => task-sdk/tests/task_sdk/log}/test_log.py (100%) delete mode 100644 task_sdk/tests/execution_time/__init__.py diff --git a/.dockerignore b/.dockerignore index dce962375a92f..40954c20ec5f2 100644 --- a/.dockerignore +++ b/.dockerignore @@ -35,7 +35,7 @@ !docs !licenses !providers/ -!task_sdk/ +!task-sdk/ # Add those folders to the context so that they are available in the CI container !scripts diff --git a/.github/boring-cyborg.yml b/.github/boring-cyborg.yml index 3c7e223209c8f..964e45a00c090 100644 --- a/.github/boring-cyborg.yml +++ b/.github/boring-cyborg.yml @@ -414,7 +414,7 @@ labelPRBasedOnFilePath: - tests/system/**/* area:task-sdk: - - task_sdk/**/* + - task-sdk/**/* area:db-migrations: - airflow/migrations/versions/* diff --git a/.github/dependabot.yml b/.github/dependabot.yml index b243de7eda35f..3806c65061895 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -22,7 +22,7 @@ updates: - /clients/python - /dev/breeze - /docker_tests - - /task_sdk + - /task-sdk - / schedule: interval: daily diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index baac8947ae294..6fd607e835ca5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -676,7 +676,7 @@ repos: name: Sync template context variable refs language: python entry: ./scripts/ci/pre_commit/template_context_key_sync.py - files: ^airflow/models/taskinstance\.py$|^task_sdk/src/airflow/sdk/definitions/context\.py$|^docs/apache-airflow/templates-ref\.rst$ + files: ^airflow/models/taskinstance\.py$|^task-sdk/src/airflow/sdk/definitions/context\.py$|^docs/apache-airflow/templates-ref\.rst$ - id: check-base-operator-usage language: pygrep name: Check BaseOperator core imports @@ -1174,9 +1174,9 @@ repos: ^tests/ | ^providers/tests/ | ^providers/.*/tests/ | - ^task_sdk/src/airflow/sdk/definitions/dag.py$ | - ^task_sdk/src/airflow/sdk/execution_time/secrets_masker.py$ | - ^task_sdk/src/airflow/sdk/definitions/_internal/node.py$ | + ^task-sdk/src/airflow/sdk/definitions/dag.py$ | + ^task-sdk/src/airflow/sdk/execution_time/secrets_masker.py$ | + ^task-sdk/src/airflow/sdk/definitions/_internal/node.py$ | ^dev/.*\.py$ | ^scripts/.*\.py$ | ^docker_tests/.*$ | @@ -1293,7 +1293,7 @@ repos: provider_packages| providers/| scripts| - task_sdk/| + task-sdk/| tests/dags/test_imports\.py ) require_serial: true @@ -1344,14 +1344,14 @@ repos: name: Run mypy for Task SDK language: python entry: ./scripts/ci/pre_commit/mypy.py --namespace-packages - files: ^task_sdk/src/airflow/sdk/.*\.py$|^task_sdk/tests//.*\.py$ + files: ^task-sdk/src/airflow/sdk/.*\.py$|^task-sdk/tests//.*\.py$ require_serial: true additional_dependencies: ['rich>=12.4.4'] - id: mypy-task-sdk stages: ['manual'] name: Run mypy for Task SDK (manual) language: python - entry: ./scripts/ci/pre_commit/mypy_folder.py task_sdk/src/airflow/sdk + entry: ./scripts/ci/pre_commit/mypy_folder.py task-sdk/src/airflow/sdk pass_filenames: false files: ^.*\.py$ require_serial: true @@ -1398,7 +1398,7 @@ repos: - id: generate-tasksdk-datamodels name: Generate Datamodels for TaskSDK client language: python - entry: uv run --active --group codegen --project apache-airflow-task-sdk --directory task_sdk -s dev/generate_models.py + entry: uv run --active --group codegen --project apache-airflow-task-sdk --directory task-sdk -s dev/generate_models.py pass_filenames: false files: ^airflow/api_fastapi/execution_api/.*\.py$ require_serial: true diff --git a/Dockerfile b/Dockerfile index f9f3b16753e4f..7cdfaec8aa211 100644 --- a/Dockerfile +++ b/Dockerfile @@ -815,7 +815,7 @@ function install_airflow() { local installation_command_flags if [[ ${AIRFLOW_INSTALLATION_METHOD} == "." ]]; then # When installing from sources - we always use `--editable` mode - installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task_sdk --editable ./devel-common" + installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task-sdk --editable ./devel-common" while IFS= read -r -d '' pyproject_toml_file; do project_folder=$(dirname ${pyproject_toml_file}) installation_command_flags="${installation_command_flags} --editable ${project_folder}" diff --git a/Dockerfile.ci b/Dockerfile.ci index f632d40dbd390..dc63936aa757d 100644 --- a/Dockerfile.ci +++ b/Dockerfile.ci @@ -585,7 +585,7 @@ function install_airflow() { local installation_command_flags if [[ ${AIRFLOW_INSTALLATION_METHOD} == "." ]]; then # When installing from sources - we always use `--editable` mode - installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task_sdk --editable ./devel-common" + installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task-sdk --editable ./devel-common" while IFS= read -r -d '' pyproject_toml_file; do project_folder=$(dirname ${pyproject_toml_file}) installation_command_flags="${installation_command_flags} --editable ${project_folder}" @@ -1001,7 +1001,7 @@ function check_force_lowest_dependencies() { echo fi set -x - uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task_sdk" --editable "./devel-common" + uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task-sdk" --editable "./devel-common" set +x } diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index d2f0d46e0c45d..fbf3626f4916c 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -968,7 +968,7 @@ def get_triggering_events() -> dict[str, list[AssetEvent]]: return triggering_events # NOTE: If you add to this dict, make sure to also update the following: - # * Context in task_sdk/src/airflow/sdk/definitions/context.py + # * Context in task-sdk/src/airflow/sdk/definitions/context.py # * KNOWN_CONTEXT_KEYS in airflow/utils/context.py # * Table in docs/apache-airflow/templates-ref.rst diff --git a/airflow/utils/context.py b/airflow/utils/context.py index 8ed44899bcb7f..507211492ebf1 100644 --- a/airflow/utils/context.py +++ b/airflow/utils/context.py @@ -47,7 +47,7 @@ from airflow.sdk.types import OutletEventAccessorsProtocol # NOTE: Please keep this in sync with the following: -# * Context in task_sdk/src/airflow/sdk/definitions/context.py +# * Context in task-sdk/src/airflow/sdk/definitions/context.py # * Table in docs/apache-airflow/templates-ref.rst KNOWN_CONTEXT_KEYS: set[str] = { "conn", diff --git a/contributing-docs/07_local_virtualenv.rst b/contributing-docs/07_local_virtualenv.rst index 177f2d433216a..18485dd57ab57 100644 --- a/contributing-docs/07_local_virtualenv.rst +++ b/contributing-docs/07_local_virtualenv.rst @@ -227,7 +227,7 @@ dependencies. For example, to install Amazon provider you need to install ``amaz .. code:: bash - pip install -e "./task_sdk" + pip install -e "./task-sdk" pip install -e "./devel-common" pip install -e "./providers/amazon" pip install -e ".[amazon]" @@ -284,7 +284,7 @@ You can run the following command in the venv that you have installed airflow in .. code:: bash pip install -e ".[google]" - pip install -e "./task_sdk" + pip install -e "./task-sdk" pip install -e "./devel-common" pip install -e "./providers/google" @@ -294,7 +294,7 @@ dependency groups will be implemented in ``pip`` - April 2025) - it will not be when installing airflow - currently with ``pip`` it is the only way to install development dependencies of the provider and is a bit convoluted. -The second installs ``task_sdk`` project - where APIs for providers are kept. +The second installs ``task-sdk`` project - where APIs for providers are kept. The third one installs google provider source code in development mode, so that modifications to the code are automatically reflected in your installed virtualenv. diff --git a/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst b/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst index 4343497329d98..cd40d8310a5b5 100644 --- a/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst +++ b/contributing-docs/quick-start-ide/contributors_quick_start_pycharm.rst @@ -39,10 +39,10 @@ Setup your project alt="Cloning github fork to Pycharm"> -3. Configure the source root directories well as for ``task_sdk`` and ``devel-common``. +3. Configure the source root directories well as for ``task-sdk`` and ``devel-common``. You also have to set "source" and "tests" root directories for each provider you want to develop (!). - This is important in Airflow 3.0 we split ``task_sdk``, ``devel-common`` and each provider to be separate + This is important in Airflow 3.0 we split ``task-sdk``, ``devel-common`` and each provider to be separate distribution - each with separate ``pyproject.toml`` file, so you need to separately add "src" and "tests" directories for each provider you develop to be respectively "source roots" and "test roots". @@ -57,12 +57,12 @@ Setup your project alt="Adding Source Root directories to Pycharm"> - You also need to add ``task_sdk`` sources (and ``devel-common`` in similar way). + You also need to add ``task-sdk`` sources (and ``devel-common`` in similar way). .. raw:: html
- Adding Source Root directories to Pycharm
diff --git a/dev/breeze/src/airflow_breeze/commands/release_management_commands.py b/dev/breeze/src/airflow_breeze/commands/release_management_commands.py index 461c5592d012e..37b33a24e7168 100644 --- a/dev/breeze/src/airflow_breeze/commands/release_management_commands.py +++ b/dev/breeze/src/airflow_breeze/commands/release_management_commands.py @@ -627,7 +627,7 @@ def prepare_airflow_packages( get_console().print("[success]Successfully prepared Airflow packages") -TASK_SDK_DIR_PATH = AIRFLOW_SOURCES_ROOT / "task_sdk" +TASK_SDK_DIR_PATH = AIRFLOW_SOURCES_ROOT / "task-sdk" TASK_SDK_DIST_DIR_PATH = TASK_SDK_DIR_PATH / "dist" @@ -692,7 +692,7 @@ def _build_package_with_docker(package_format: str): "-e", "GITHUB_ACTIONS", "-w", - "/opt/airflow/task_sdk", + "/opt/airflow/task-sdk", AIRFLOW_BUILD_IMAGE_TAG, "bash", "-c", @@ -707,7 +707,7 @@ def _build_package_with_docker(package_format: str): DIST_DIR.mkdir(parents=True, exist_ok=True) get_console().print() # Copy all files in the dist directory in container to the host dist directory (note '/.' in SRC) - run_command(["docker", "cp", f"{container_id}:/opt/airflow/task_sdk/dist/.", "./dist"], check=True) + run_command(["docker", "cp", f"{container_id}:/opt/airflow/task-sdk/dist/.", "./dist"], check=True) run_command(["docker", "rm", "--force", container_id], check=False, stdout=DEVNULL, stderr=DEVNULL) if use_local_hatch: diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands.py b/dev/breeze/src/airflow_breeze/commands/testing_commands.py index 83ebea2c96942..5939de725ed33 100644 --- a/dev/breeze/src/airflow_breeze/commands/testing_commands.py +++ b/dev/breeze/src/airflow_breeze/commands/testing_commands.py @@ -154,7 +154,7 @@ def docker_compose_tests( sys.exit(return_code) -TEST_PROGRESS_REGEXP = r"tests/.*|providers/.*/tests/.*|task_sdk/tests/.*|.*=====.*" +TEST_PROGRESS_REGEXP = r"tests/.*|providers/.*/tests/.*|task-sdk/tests/.*|.*=====.*" PERCENT_TEST_PROGRESS_REGEXP = r"^tests/.*\[[ \d%]*\].*|^\..*\[[ \d%]*\].*" diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index 9c64c567c413a..fbbd1b43bb354 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -287,7 +287,7 @@ def all_task_sdk_test_packages() -> list[str]: return sorted( [ candidate.name - for candidate in (AIRFLOW_SOURCES_ROOT / "task_sdk" / "tests").iterdir() + for candidate in (AIRFLOW_SOURCES_ROOT / "task-sdk" / "tests").iterdir() if candidate.is_dir() and candidate.name != "__pycache__" ] ) diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py index 289466c7aaa8b..0c280698569b2 100644 --- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py +++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py @@ -92,7 +92,7 @@ ("hooks", "/opt/airflow/hooks"), ("logs", "/root/airflow/logs"), ("providers", "/opt/airflow/providers"), - ("task_sdk", "/opt/airflow/task_sdk"), + ("task-sdk", "/opt/airflow/task-sdk"), ("pyproject.toml", "/opt/airflow/pyproject.toml"), ("scripts", "/opt/airflow/scripts"), ("scripts/docker/entrypoint_ci.sh", "/entrypoint"), diff --git a/dev/breeze/src/airflow_breeze/utils/run_tests.py b/dev/breeze/src/airflow_breeze/utils/run_tests.py index 174ad19975a07..814e1c91eacde 100644 --- a/dev/breeze/src/airflow_breeze/utils/run_tests.py +++ b/dev/breeze/src/airflow_breeze/utils/run_tests.py @@ -168,7 +168,7 @@ def get_excluded_provider_args(python_version: str) -> list[str]: "Serialization": [ "tests/serialization", ], - "TaskSDK": ["task_sdk/tests"], + "TaskSDK": ["task-sdk/tests"], "WWW": [ "tests/www", ], @@ -200,7 +200,7 @@ def get_excluded_provider_args(python_version: str) -> list[str]: TEST_GROUP_TO_TEST_FOLDERS: dict[GroupOfTests, list[str]] = { GroupOfTests.CORE: ["tests"], GroupOfTests.PROVIDERS: ALL_PROVIDER_TEST_FOLDERS, - GroupOfTests.TASK_SDK: ["task_sdk/tests"], + GroupOfTests.TASK_SDK: ["task-sdk/tests"], GroupOfTests.HELM: ["helm_tests"], GroupOfTests.INTEGRATION_CORE: ["tests/integration"], GroupOfTests.INTEGRATION_PROVIDERS: ALL_PROVIDER_INTEGRATION_TEST_FOLDERS, diff --git a/dev/breeze/src/airflow_breeze/utils/selective_checks.py b/dev/breeze/src/airflow_breeze/utils/selective_checks.py index b4a0d2c4310b8..7d7be448088d3 100644 --- a/dev/breeze/src/airflow_breeze/utils/selective_checks.py +++ b/dev/breeze/src/airflow_breeze/utils/selective_checks.py @@ -190,7 +190,7 @@ def __hash__(self): r"^chart", r"^providers/.*/src/", r"^providers/.*/docs/", - r"^task_sdk/src/", + r"^task-sdk/src/", r"^tests/system", r"^CHANGELOG\.txt", r"^airflow/config_templates/config\.yml", @@ -226,8 +226,8 @@ def __hash__(self): r"^providers/tests/", r"^providers/.*/src/", r"^providers/.*/tests/", - r"^task_sdk/src/", - r"^task_sdk/tests/", + r"^task-sdk/src/", + r"^task-sdk/tests/", r"^tests", r"^devel-common", r"^kubernetes_tests", @@ -246,18 +246,18 @@ def __hash__(self): r"^devel-common/.*\.py$", ], FileGroupForCi.TASK_SDK_FILES: [ - r"^task_sdk/src/airflow/sdk/.*\.py$", - r"^task_sdk/tests/.*\.py$", + r"^task-sdk/src/airflow/sdk/.*\.py$", + r"^task-sdk/tests/.*\.py$", ], FileGroupForCi.ASSET_FILES: [ r"^airflow/assets/", r"^airflow/models/assets/", - r"^task_sdk/src/airflow/sdk/definitions/asset/", + r"^task-sdk/src/airflow/sdk/definitions/asset/", r"^airflow/datasets/", ], FileGroupForCi.UNIT_TEST_FILES: [ r"^tests/", - r"^task_sdk/tests/", + r"^task-sdk/tests/", r"^providers/.*/tests/", r"^dev/breeze/tests/", ], @@ -274,8 +274,8 @@ def __hash__(self): r"^docs/.*", r"^providers/.*/tests/.*", r"^tests/dags/test_imports.py", - r"^task_sdk/src/airflow/sdk/.*\.py$", - r"^task_sdk/tests/.*\.py$", + r"^task-sdk/src/airflow/sdk/.*\.py$", + r"^task-sdk/tests/.*\.py$", ] } ) @@ -307,8 +307,8 @@ def __hash__(self): r"^providers/.*/tests/", ], SelectiveTaskSdkTestType.TASK_SDK: [ - r"^task_sdk/src/", - r"^task_sdk/tests/", + r"^task-sdk/src/", + r"^task-sdk/tests/", ], SelectiveCoreTestType.SERIALIZATION: [ r"^airflow/serialization/", diff --git a/dev/breeze/tests/test_pytest_args_for_test_types.py b/dev/breeze/tests/test_pytest_args_for_test_types.py index 33db4c4974883..4500768a568ea 100644 --- a/dev/breeze/tests/test_pytest_args_for_test_types.py +++ b/dev/breeze/tests/test_pytest_args_for_test_types.py @@ -304,7 +304,7 @@ def test_pytest_args_for_missing_provider(): GroupOfTests.TASK_SDK, "All", [ - "task_sdk/tests", + "task-sdk/tests", ], ), ( diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py index fd2afabf6ae31..250434548a9af 100644 --- a/dev/breeze/tests/test_selective_checks.py +++ b/dev/breeze/tests/test_selective_checks.py @@ -492,7 +492,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): ), ( pytest.param( - ("task_sdk/src/airflow/sdk/random.py",), + ("task-sdk/src/airflow/sdk/random.py",), { "all-python-versions": "['3.9']", "all-python-versions-list-as-string": "3.9", @@ -1688,7 +1688,7 @@ def test_expected_output_push( ( "airflow/assets/", "airflow/models/assets/", - "task_sdk/src/airflow/sdk/definitions/asset/", + "task-sdk/src/airflow/sdk/definitions/asset/", "airflow/datasets/", ), { @@ -2455,7 +2455,7 @@ def test_provider_compatibility_checks(labels: tuple[str, ...], expected_outputs id="Airflow mypy checks on airflow files with model changes.", ), pytest.param( - ("task_sdk/src/airflow/sdk/a_file.py",), + ("task-sdk/src/airflow/sdk/a_file.py",), { "needs-mypy": "true", "mypy-checks": "['mypy-providers', 'mypy-task-sdk']", @@ -2594,7 +2594,7 @@ def test_pr_labels( id="Caplog is in the git diff Providers", ), pytest.param( - ("task_sdk/tests/definitions/test_dag.py",), + ("task-sdk/tests/definitions/test_dag.py",), (), GithubEvents.PULL_REQUEST, id="Caplog is in the git diff TaskSDK", @@ -2651,7 +2651,7 @@ def test_is_log_mocked_in_the_tests_fail( id="Caplog is in the git diff Providers", ), pytest.param( - ("task_sdk/tests/definitions/test_dag.py",), + ("task-sdk/tests/definitions/test_dag.py",), (), GithubEvents.PULL_REQUEST, id="Caplog is in the git diff TaskSDK", @@ -2711,7 +2711,7 @@ def test_is_log_mocked_in_the_tests_fail_formatted( id="Caplog is in the git diff Providers", ), pytest.param( - ("task_sdk/tests/definitions/test_dag.py",), + ("task-sdk/tests/definitions/test_dag.py",), (), GithubEvents.PULL_REQUEST, id="Caplog is in the git diff TaskSDK", @@ -2759,7 +2759,7 @@ def test_is_log_mocked_in_the_tests_not_fail( id="Caplog is in the git diff Providers", ), pytest.param( - ("task_sdk/tests/definitions/test_dag.py",), + ("task-sdk/tests/definitions/test_dag.py",), (LOG_WITHOUT_MOCK_IN_TESTS_EXCEPTION_LABEL,), GithubEvents.PULL_REQUEST, id="Caplog is in the git diff TaskSDK", diff --git a/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst b/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst index 7d727bc01c421..417443192d0b0 100644 --- a/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst +++ b/docs/apache-airflow/administration-and-deployment/logging-monitoring/callbacks.rst @@ -97,4 +97,4 @@ In the following example, failures in any task call the ``task_failure_alert`` f to be executed in the desired event. Simply pass a list of callback functions to the callback args when defining your DAG/task callbacks: e.g ``on_failure_callback=[callback_func_1, callback_func_2]`` -Full list of variables available in ``context`` in :doc:`docs <../../templates-ref>` and `code `_. +Full list of variables available in ``context`` in :doc:`docs <../../templates-ref>` and `code `_. diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py index 5ad2e893c1aaa..256bc133f1a41 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_manager.py @@ -20,7 +20,7 @@ import logging import tempfile from datetime import datetime -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol from unittest import mock from unittest.mock import MagicMock @@ -46,12 +46,16 @@ if TYPE_CHECKING: try: + from airflow.sdk.api.datamodels._generated import TIRunContext from airflow.sdk.definitions.context import Context + except ImportError: # TODO: Remove once provider drops support for Airflow 2 + # TIRunContext is only used in Airflow 3 tests from airflow.utils.context import Context - from task_sdk.tests.conftest import MakeTIContextCallable + TIRunContext = Any # type: ignore[misc, assignment] + if AIRFLOW_V_2_10_PLUS: @@ -414,6 +418,24 @@ def set_dag(what: StartupDetails, dag_id: str, task: BaseOperator) -> RuntimeTas return set_dag +class MakeTIContextCallable(Protocol): + def __call__( + self, + dag_id: str = ..., + run_id: str = ..., + logical_date: str | datetime = ..., + data_interval_start: str | datetime = ..., + data_interval_end: str | datetime = ..., + clear_number: int = ..., + start_date: str | datetime = ..., + run_after: str | datetime = ..., + run_type: str = ..., + task_reschedule_count: int = ..., + conf: dict[str, Any] | None = ..., + ) -> TIRunContext: ... + + +# Only needed in Airflow 3 @pytest.fixture def make_ti_context() -> MakeTIContextCallable: """Factory for creating TIRunContext objects.""" diff --git a/pyproject.toml b/pyproject.toml index 12028e1908e8a..0b11eb5163c0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -294,7 +294,7 @@ section-order = [ # Make sure we put the "dev" imports at the end, not as a third-party module [tool.ruff.lint.isort.sections] -testing = ["dev", "providers.tests", "task_sdk.tests", "tests_common", "tests"] +testing = ["dev", "providers.tests", "tests_common", "tests"] [tool.ruff.lint.extend-per-file-ignores] "airflow/__init__.py" = ["F401", "TC004", "I002"] @@ -916,6 +916,6 @@ members = [ "providers/ydb", "providers/zendesk", ".", - "task_sdk", + "task-sdk", "devel-common", ] diff --git a/scripts/ci/docker-compose/local.yml b/scripts/ci/docker-compose/local.yml index f23a2a1f9c896..156f20a071c3f 100644 --- a/scripts/ci/docker-compose/local.yml +++ b/scripts/ci/docker-compose/local.yml @@ -89,8 +89,8 @@ services: source: ../../../providers target: /opt/airflow/providers - type: bind - source: ../../../task_sdk - target: /opt/airflow/task_sdk + source: ../../../task-sdk + target: /opt/airflow/task-sdk - type: bind source: ../../../pyproject.toml target: /opt/airflow/pyproject.toml diff --git a/scripts/ci/docker-compose/providers-and-tests-sources.yml b/scripts/ci/docker-compose/providers-and-tests-sources.yml index 91031e22cceae..2f5eff3835e31 100644 --- a/scripts/ci/docker-compose/providers-and-tests-sources.yml +++ b/scripts/ci/docker-compose/providers-and-tests-sources.yml @@ -30,10 +30,10 @@ services: - ../../../empty:/opt/airflow/airflow # But keep airflow/__init__.py to make hatchling happy when building empty package - ../../../airflow/__init__.py:/opt/airflow/airflow/__init__.py - # Remove task_sdk sources from container - - ../../../empty:/opt/airflow/task_sdk/src - # But keep task_sdk/a__init__.py to make hatchling happy when building empty package - - ../../../task_sdk/src/airflow/sdk/__init__.py:/opt/airflow/task_sdk/src/airflow/sdk/__init__.py + # Remove task-sdk sources from container + - ../../../empty:/opt/airflow/task-sdk/src + # But keep task-sdk/a__init__.py to make hatchling happy when building empty package + - ../../../task-sdk/src/airflow/sdk/__init__.py:/opt/airflow/task-sdk/src/airflow/sdk/__init__.py # And keep tests - ../../../tests/:/opt/airflow/tests:cached # Mount providers to make sure that we have the latest providers - both tests and sources diff --git a/scripts/ci/docker-compose/remove-sources.yml b/scripts/ci/docker-compose/remove-sources.yml index 5c31f0baee6f3..fd7b73c93dff1 100644 --- a/scripts/ci/docker-compose/remove-sources.yml +++ b/scripts/ci/docker-compose/remove-sources.yml @@ -22,8 +22,8 @@ services: volumes: # Removes airflow sources from container - ../../../empty:/opt/airflow/airflow - # Removes task_sdk sources from container - - ../../../empty:/opt/airflow/task_sdk/src + # Removes task-sdk sources from container + - ../../../empty:/opt/airflow/task-sdk/src # Remove all provider sources from container # START automatically generated volumes by generate-volumes-for-sources pre-commit - ../../../empty:/opt/airflow/providers/airbyte/src diff --git a/scripts/ci/kubernetes/k8s_requirements.txt b/scripts/ci/kubernetes/k8s_requirements.txt index 03908cd6818ef..0fbbf3827c3a0 100644 --- a/scripts/ci/kubernetes/k8s_requirements.txt +++ b/scripts/ci/kubernetes/k8s_requirements.txt @@ -1,4 +1,4 @@ -e ./providers/standard --e ./task_sdk +-e ./task-sdk -e ./devel-common -e ./providers/cncf/kubernetes diff --git a/scripts/ci/pre_commit/base_operator_partial_arguments.py b/scripts/ci/pre_commit/base_operator_partial_arguments.py index 070ffe767cccc..7b3b81bb838c4 100755 --- a/scripts/ci/pre_commit/base_operator_partial_arguments.py +++ b/scripts/ci/pre_commit/base_operator_partial_arguments.py @@ -27,9 +27,9 @@ ROOT_DIR = pathlib.Path(__file__).resolve().parents[3] BASEOPERATOR_PY = ROOT_DIR.joinpath("airflow", "models", "baseoperator.py") -SDK_BASEOPERATOR_PY = ROOT_DIR.joinpath("task_sdk", "src", "airflow", "sdk", "definitions", "baseoperator.py") +SDK_BASEOPERATOR_PY = ROOT_DIR.joinpath("task-sdk", "src", "airflow", "sdk", "definitions", "baseoperator.py") SDK_MAPPEDOPERATOR_PY = ROOT_DIR.joinpath( - "task_sdk", "src", "airflow", "sdk", "definitions", "mappedoperator.py" + "task-sdk", "src", "airflow", "sdk", "definitions", "mappedoperator.py" ) IGNORED = { diff --git a/scripts/ci/pre_commit/mypy_folder.py b/scripts/ci/pre_commit/mypy_folder.py index 09a643d45ccbd..ee4c35b076163 100755 --- a/scripts/ci/pre_commit/mypy_folder.py +++ b/scripts/ci/pre_commit/mypy_folder.py @@ -38,7 +38,7 @@ *[f"providers/{provider_id.replace('.', '/')}/src" for provider_id in get_all_provider_ids()], "dev", "docs", - "task_sdk/src/airflow/sdk", + "task-sdk/src/airflow/sdk", "all_providers", ] @@ -69,8 +69,8 @@ elif mypy_folder.startswith("providers/"): arguments.append(f"{Path(mypy_folder).parent.as_posix()}/tests") namespace_packages = True - if mypy_folder == "task_sdk/src/airflow/sdk": - arguments.append("task_sdk/tests") + if mypy_folder == "task-sdk/src/airflow/sdk": + arguments.append("task-sdk/tests") namespace_packages = True if mypy_folder == "airflow": arguments.append("tests") diff --git a/scripts/ci/pre_commit/sync_init_decorator.py b/scripts/ci/pre_commit/sync_init_decorator.py index 7b02136ead318..18efca02dd519 100755 --- a/scripts/ci/pre_commit/sync_init_decorator.py +++ b/scripts/ci/pre_commit/sync_init_decorator.py @@ -28,7 +28,7 @@ ROOT = pathlib.Path(__file__).resolve().parents[3] PACKAGE_ROOT = ROOT.joinpath("airflow") -SDK_DEFINITIONS_PKG = ROOT.joinpath("task_sdk", "src", "airflow", "sdk", "definitions") +SDK_DEFINITIONS_PKG = ROOT.joinpath("task-sdk", "src", "airflow", "sdk", "definitions") DAG_PY = SDK_DEFINITIONS_PKG.joinpath("dag.py") TG_PY = SDK_DEFINITIONS_PKG.joinpath("taskgroup.py") DECOS_TG_PY = PACKAGE_ROOT.joinpath("decorators", "task_group.py") diff --git a/scripts/ci/pre_commit/template_context_key_sync.py b/scripts/ci/pre_commit/template_context_key_sync.py index 270a3a695b263..f29f3a1a9c7af 100755 --- a/scripts/ci/pre_commit/template_context_key_sync.py +++ b/scripts/ci/pre_commit/template_context_key_sync.py @@ -27,9 +27,9 @@ ROOT_DIR = pathlib.Path(__file__).resolve().parents[3] -TASKRUNNER_PY = ROOT_DIR.joinpath("task_sdk", "src", "airflow", "sdk", "execution_time", "task_runner.py") +TASKRUNNER_PY = ROOT_DIR.joinpath("task-sdk", "src", "airflow", "sdk", "execution_time", "task_runner.py") CONTEXT_PY = ROOT_DIR.joinpath("airflow", "utils", "context.py") -CONTEXT_HINT = ROOT_DIR.joinpath("task_sdk", "src", "airflow", "sdk", "definitions", "context.py") +CONTEXT_HINT = ROOT_DIR.joinpath("task-sdk", "src", "airflow", "sdk", "definitions", "context.py") TEMPLATES_REF_RST = ROOT_DIR.joinpath("docs", "apache-airflow", "templates-ref.rst") # These are only conditionally set diff --git a/scripts/ci/testing/run_unit_tests.sh b/scripts/ci/testing/run_unit_tests.sh index b602237ba7e52..cace8a1dfbfc1 100755 --- a/scripts/ci/testing/run_unit_tests.sh +++ b/scripts/ci/testing/run_unit_tests.sh @@ -135,7 +135,7 @@ function run_tests() { core_tests elif [[ "${TEST_GROUP}" == "providers" ]]; then providers_tests - elif [[ "${TEST_GROUP}" == "task_sdk" ]]; then + elif [[ "${TEST_GROUP}" == "task-sdk" ]]; then task_sdk_tests else echo "Unknown test group: ${TEST_GROUP}" diff --git a/scripts/docker/entrypoint_ci.sh b/scripts/docker/entrypoint_ci.sh index e6abb50f73323..4caa698e24256 100755 --- a/scripts/docker/entrypoint_ci.sh +++ b/scripts/docker/entrypoint_ci.sh @@ -331,7 +331,7 @@ function check_force_lowest_dependencies() { echo fi set -x - uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task_sdk" --editable "./devel-common" + uv pip install --python "$(which python)" --resolution lowest-direct --upgrade --editable ".${EXTRA}" --editable "./task-sdk" --editable "./devel-common" set +x } diff --git a/scripts/docker/install_airflow.sh b/scripts/docker/install_airflow.sh index d34ce3124ba36..5511c0d3afdef 100644 --- a/scripts/docker/install_airflow.sh +++ b/scripts/docker/install_airflow.sh @@ -49,7 +49,7 @@ function install_airflow() { local installation_command_flags if [[ ${AIRFLOW_INSTALLATION_METHOD} == "." ]]; then # When installing from sources - we always use `--editable` mode - installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task_sdk --editable ./devel-common" + installation_command_flags="--editable .[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION} --editable ./task-sdk --editable ./devel-common" while IFS= read -r -d '' pyproject_toml_file; do project_folder=$(dirname ${pyproject_toml_file}) installation_command_flags="${installation_command_flags} --editable ${project_folder}" diff --git a/scripts/in_container/install_airflow_and_providers.py b/scripts/in_container/install_airflow_and_providers.py index 6892d99651e3b..33032a011c777 100755 --- a/scripts/in_container/install_airflow_and_providers.py +++ b/scripts/in_container/install_airflow_and_providers.py @@ -249,7 +249,7 @@ def find_installation_spec( ) if airflow_extras: airflow_package_spec += airflow_extras - # We always install latest task_sdk - it's independent from Airflow + # We always install latest task-sdk - it's independent from Airflow airflow_task_sdk_spec = find_airflow_task_sdk_package(extension) if airflow_task_sdk_spec: airflow_task_sdk_constraints_location = get_airflow_constraints_location( @@ -287,7 +287,7 @@ def find_installation_spec( github_repository=github_repository, python_version=python_version, ) - console.print(f"\nInstalling airflow task_sdk from remote spec {use_airflow_version}\n") + console.print(f"\nInstalling airflow task-sdk from remote spec {use_airflow_version}\n") airflow_task_sdk_package = f"apache-airflow-task-sdk @ {use_airflow_version}" airflow_constraints_location = get_airflow_constraints_location( airflow_skip_constraints=airflow_skip_constraints, @@ -318,7 +318,7 @@ def find_installation_spec( python_version=python_version, ) console.print( - "\nDo not install airflow task_sdk. It should be installed automatically if needed " + "\nDo not install airflow task-sdk. It should be installed automatically if needed " "by providers." ) airflow_task_sdk_package = None diff --git a/task_sdk/README.md b/task-sdk/README.md similarity index 100% rename from task_sdk/README.md rename to task-sdk/README.md diff --git a/task_sdk/dev/generate_models.py b/task-sdk/dev/generate_models.py similarity index 100% rename from task_sdk/dev/generate_models.py rename to task-sdk/dev/generate_models.py diff --git a/task_sdk/pyproject.toml b/task-sdk/pyproject.toml similarity index 80% rename from task_sdk/pyproject.toml rename to task-sdk/pyproject.toml index d980c0aae8259..8f983c3e6b5c1 100644 --- a/task_sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -141,3 +141,42 @@ custom-formatters = ['dev.datamodel_code_formatter',] url = 'http://0.0.0.0:8080/execution/openapi.json' output = 'src/airflow/sdk/api/datamodels/_generated.py' + +## pytest settings ## +[tool.pytest.ini_options] +addopts = [ + "--tb=short", + "-rasl", + "--verbosity=2", + # Disable `flaky` plugin for pytest. This plugin conflicts with `rerunfailures` because provide the same marker. + "-p", "no:flaky", + # Disable `nose` builtin plugin for pytest. This feature is deprecated in 7.2 and will be removed in pytest>=8 + "-p", "no:nose", + # Disable support of a legacy `LocalPath` in favor of stdlib `pathlib.Path`. + "-p", "no:legacypath", + # Disable warnings summary, because we use our warning summary. + "--disable-warnings", + "--asyncio-mode=strict", +] + +norecursedirs = [ + ".eggs", +] +log_level = "INFO" +filterwarnings = [ + "error::pytest.PytestCollectionWarning", + "error::pytest.PytestReturnNotNoneWarning", +] +python_files = [ + "test_*.py", +] +testpaths = [ + "tests", +] +asyncio_default_fixture_loop_scope = "function" + +pythonpath = "tests" + +# Keep temporary directories (created by `tmp_path`) for 2 recent runs only failed tests. +tmp_path_retention_count = "2" +tmp_path_retention_policy = "failed" diff --git a/task_sdk/src/airflow/__init__.py b/task-sdk/src/airflow/__init__.py similarity index 100% rename from task_sdk/src/airflow/__init__.py rename to task-sdk/src/airflow/__init__.py diff --git a/task_sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py similarity index 100% rename from task_sdk/src/airflow/sdk/__init__.py rename to task-sdk/src/airflow/sdk/__init__.py diff --git a/task_sdk/__init__.py b/task-sdk/src/airflow/sdk/api/__init__.py similarity index 100% rename from task_sdk/__init__.py rename to task-sdk/src/airflow/sdk/api/__init__.py diff --git a/task_sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py similarity index 100% rename from task_sdk/src/airflow/sdk/api/client.py rename to task-sdk/src/airflow/sdk/api/client.py diff --git a/task_sdk/src/airflow/sdk/api/__init__.py b/task-sdk/src/airflow/sdk/api/datamodels/__init__.py similarity index 100% rename from task_sdk/src/airflow/sdk/api/__init__.py rename to task-sdk/src/airflow/sdk/api/datamodels/__init__.py diff --git a/task_sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py similarity index 100% rename from task_sdk/src/airflow/sdk/api/datamodels/_generated.py rename to task-sdk/src/airflow/sdk/api/datamodels/_generated.py diff --git a/task_sdk/src/airflow/sdk/api/datamodels/activities.py b/task-sdk/src/airflow/sdk/api/datamodels/activities.py similarity index 100% rename from task_sdk/src/airflow/sdk/api/datamodels/activities.py rename to task-sdk/src/airflow/sdk/api/datamodels/activities.py diff --git a/task_sdk/src/airflow/sdk/api/datamodels/__init__.py b/task-sdk/src/airflow/sdk/definitions/__init__.py similarity index 100% rename from task_sdk/src/airflow/sdk/api/datamodels/__init__.py rename to task-sdk/src/airflow/sdk/definitions/__init__.py diff --git a/task_sdk/src/airflow/sdk/definitions/__init__.py b/task-sdk/src/airflow/sdk/definitions/_internal/__init__.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/__init__.py rename to task-sdk/src/airflow/sdk/definitions/_internal/__init__.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py rename to task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/contextmanager.py b/task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/contextmanager.py rename to task-sdk/src/airflow/sdk/definitions/_internal/contextmanager.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py b/task-sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py rename to task-sdk/src/airflow/sdk/definitions/_internal/dag_parsing_context.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/decorators.py b/task-sdk/src/airflow/sdk/definitions/_internal/decorators.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/decorators.py rename to task-sdk/src/airflow/sdk/definitions/_internal/decorators.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/expandinput.py b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/expandinput.py rename to task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/mixins.py b/task-sdk/src/airflow/sdk/definitions/_internal/mixins.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/mixins.py rename to task-sdk/src/airflow/sdk/definitions/_internal/mixins.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/node.py b/task-sdk/src/airflow/sdk/definitions/_internal/node.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/node.py rename to task-sdk/src/airflow/sdk/definitions/_internal/node.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/templater.py b/task-sdk/src/airflow/sdk/definitions/_internal/templater.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/templater.py rename to task-sdk/src/airflow/sdk/definitions/_internal/templater.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/types.py b/task-sdk/src/airflow/sdk/definitions/_internal/types.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/types.py rename to task-sdk/src/airflow/sdk/definitions/_internal/types.py diff --git a/task_sdk/src/airflow/sdk/definitions/asset/__init__.py b/task-sdk/src/airflow/sdk/definitions/asset/__init__.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/asset/__init__.py rename to task-sdk/src/airflow/sdk/definitions/asset/__init__.py diff --git a/task_sdk/src/airflow/sdk/definitions/asset/decorators.py b/task-sdk/src/airflow/sdk/definitions/asset/decorators.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/asset/decorators.py rename to task-sdk/src/airflow/sdk/definitions/asset/decorators.py diff --git a/task_sdk/src/airflow/sdk/definitions/asset/metadata.py b/task-sdk/src/airflow/sdk/definitions/asset/metadata.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/asset/metadata.py rename to task-sdk/src/airflow/sdk/definitions/asset/metadata.py diff --git a/task_sdk/src/airflow/sdk/definitions/baseoperator.py b/task-sdk/src/airflow/sdk/definitions/baseoperator.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/baseoperator.py rename to task-sdk/src/airflow/sdk/definitions/baseoperator.py diff --git a/task_sdk/src/airflow/sdk/definitions/connection.py b/task-sdk/src/airflow/sdk/definitions/connection.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/connection.py rename to task-sdk/src/airflow/sdk/definitions/connection.py diff --git a/task_sdk/src/airflow/sdk/definitions/context.py b/task-sdk/src/airflow/sdk/definitions/context.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/context.py rename to task-sdk/src/airflow/sdk/definitions/context.py diff --git a/task_sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/dag.py rename to task-sdk/src/airflow/sdk/definitions/dag.py diff --git a/task_sdk/src/airflow/sdk/definitions/edges.py b/task-sdk/src/airflow/sdk/definitions/edges.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/edges.py rename to task-sdk/src/airflow/sdk/definitions/edges.py diff --git a/task_sdk/src/airflow/sdk/definitions/macros.py b/task-sdk/src/airflow/sdk/definitions/macros.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/macros.py rename to task-sdk/src/airflow/sdk/definitions/macros.py diff --git a/task_sdk/src/airflow/sdk/definitions/mappedoperator.py b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/mappedoperator.py rename to task-sdk/src/airflow/sdk/definitions/mappedoperator.py diff --git a/task_sdk/src/airflow/sdk/definitions/param.py b/task-sdk/src/airflow/sdk/definitions/param.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/param.py rename to task-sdk/src/airflow/sdk/definitions/param.py diff --git a/task_sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/taskgroup.py rename to task-sdk/src/airflow/sdk/definitions/taskgroup.py diff --git a/task_sdk/src/airflow/sdk/definitions/template.py b/task-sdk/src/airflow/sdk/definitions/template.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/template.py rename to task-sdk/src/airflow/sdk/definitions/template.py diff --git a/task_sdk/src/airflow/sdk/definitions/variable.py b/task-sdk/src/airflow/sdk/definitions/variable.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/variable.py rename to task-sdk/src/airflow/sdk/definitions/variable.py diff --git a/task_sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/xcom_arg.py rename to task-sdk/src/airflow/sdk/definitions/xcom_arg.py diff --git a/task_sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py similarity index 100% rename from task_sdk/src/airflow/sdk/exceptions.py rename to task-sdk/src/airflow/sdk/exceptions.py diff --git a/task_sdk/src/airflow/sdk/execution_time/__init__.py b/task-sdk/src/airflow/sdk/execution_time/__init__.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/__init__.py rename to task-sdk/src/airflow/sdk/execution_time/__init__.py diff --git a/task_sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/comms.py rename to task-sdk/src/airflow/sdk/execution_time/comms.py diff --git a/task_sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/context.py rename to task-sdk/src/airflow/sdk/execution_time/context.py diff --git a/task_sdk/src/airflow/sdk/execution_time/execute_workload.py b/task-sdk/src/airflow/sdk/execution_time/execute_workload.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/execute_workload.py rename to task-sdk/src/airflow/sdk/execution_time/execute_workload.py diff --git a/task_sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/lazy_sequence.py rename to task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py diff --git a/task_sdk/src/airflow/sdk/execution_time/secrets_masker.py b/task-sdk/src/airflow/sdk/execution_time/secrets_masker.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/secrets_masker.py rename to task-sdk/src/airflow/sdk/execution_time/secrets_masker.py diff --git a/task_sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/supervisor.py rename to task-sdk/src/airflow/sdk/execution_time/supervisor.py diff --git a/task_sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py similarity index 100% rename from task_sdk/src/airflow/sdk/execution_time/task_runner.py rename to task-sdk/src/airflow/sdk/execution_time/task_runner.py diff --git a/task_sdk/src/airflow/sdk/log.py b/task-sdk/src/airflow/sdk/log.py similarity index 100% rename from task_sdk/src/airflow/sdk/log.py rename to task-sdk/src/airflow/sdk/log.py diff --git a/task_sdk/src/airflow/sdk/py.typed b/task-sdk/src/airflow/sdk/py.typed similarity index 100% rename from task_sdk/src/airflow/sdk/py.typed rename to task-sdk/src/airflow/sdk/py.typed diff --git a/task_sdk/src/airflow/sdk/types.py b/task-sdk/src/airflow/sdk/types.py similarity index 100% rename from task_sdk/src/airflow/sdk/types.py rename to task-sdk/src/airflow/sdk/types.py diff --git a/task_sdk/src/airflow/sdk/definitions/_internal/__init__.py b/task-sdk/tests/__init__.py similarity index 100% rename from task_sdk/src/airflow/sdk/definitions/_internal/__init__.py rename to task-sdk/tests/__init__.py diff --git a/task_sdk/tests/conftest.py b/task-sdk/tests/conftest.py similarity index 99% rename from task_sdk/tests/conftest.py rename to task-sdk/tests/conftest.py index b075d2e73ad65..4d2560106835a 100644 --- a/task_sdk/tests/conftest.py +++ b/task-sdk/tests/conftest.py @@ -89,7 +89,7 @@ def __call__(self, _: WrappedLogger, method_name: str, event: EventDict | bytes) @pytest.fixture def test_dags_dir(): - return Path(__file__).parent.joinpath("dags") + return Path(__file__).parent.joinpath("task_sdk", "dags") @pytest.fixture diff --git a/task-sdk/tests/task_sdk/__init__.py b/task-sdk/tests/task_sdk/__init__.py new file mode 100644 index 0000000000000..059fa06ad2846 --- /dev/null +++ b/task-sdk/tests/task_sdk/__init__.py @@ -0,0 +1,45 @@ +# 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 httpx + +from airflow.sdk.api.client import Client +from airflow.sdk.execution_time.comms import BundleInfo + +FAKE_BUNDLE = BundleInfo(name="anything", version="any") + + +def make_client(transport: httpx.MockTransport) -> Client: + """Get a client with a custom transport.""" + return Client(base_url="test://server", token="", transport=transport) + + +def make_client_w_dry_run() -> Client: + """Get a client with dry_run enabled.""" + return Client(base_url=None, dry_run=True, token="") + + +def make_client_w_responses(responses: list[httpx.Response]) -> Client: + """Get a client with custom responses.""" + + def handle_request(request: httpx.Request) -> httpx.Response: + return responses.pop(0) + + return Client( + base_url=None, dry_run=True, token="", mounts={"'http://": httpx.MockTransport(handle_request)} + ) diff --git a/task_sdk/tests/__init__.py b/task-sdk/tests/task_sdk/api/__init__.py similarity index 100% rename from task_sdk/tests/__init__.py rename to task-sdk/tests/task_sdk/api/__init__.py diff --git a/task_sdk/tests/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py similarity index 97% rename from task_sdk/tests/api/test_client.py rename to task-sdk/tests/task_sdk/api/test_client.py index abadb7ef6832f..bcfe92d481c9b 100644 --- a/task_sdk/tests/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -23,8 +23,9 @@ import httpx import pytest import uuid6 +from task_sdk import make_client, make_client_w_dry_run, make_client_w_responses -from airflow.sdk.api.client import Client, RemoteValidationError, ServerResponseError +from airflow.sdk.api.client import RemoteValidationError, ServerResponseError from airflow.sdk.api.datamodels._generated import ConnectionResponse, VariableResponse, XComResponse from airflow.sdk.exceptions import ErrorType from airflow.sdk.execution_time.comms import DeferTask, ErrorResponse, RescheduleTask @@ -32,27 +33,6 @@ from airflow.utils.state import TerminalTIState -def make_client(transport: httpx.MockTransport) -> Client: - """Get a client with a custom transport""" - return Client(base_url="test://server", token="", transport=transport) - - -def make_client_w_dry_run() -> Client: - """Get a client with dry_run enabled""" - return Client(base_url=None, dry_run=True, token="") - - -def make_client_w_responses(responses: list[httpx.Response]) -> Client: - """Helper fixture to create a mock client with custom responses.""" - - def handle_request(request: httpx.Request) -> httpx.Response: - return responses.pop(0) - - return Client( - base_url=None, dry_run=True, token="", mounts={"'http://": httpx.MockTransport(handle_request)} - ) - - class TestClient: @pytest.mark.parametrize( ["path", "json_response"], diff --git a/task_sdk/tests/dags/dag_parsing_context.py b/task-sdk/tests/task_sdk/dags/dag_parsing_context.py similarity index 100% rename from task_sdk/tests/dags/dag_parsing_context.py rename to task-sdk/tests/task_sdk/dags/dag_parsing_context.py diff --git a/task_sdk/tests/dags/super_basic.py b/task-sdk/tests/task_sdk/dags/super_basic.py similarity index 100% rename from task_sdk/tests/dags/super_basic.py rename to task-sdk/tests/task_sdk/dags/super_basic.py diff --git a/task_sdk/tests/dags/super_basic_deferred_run.py b/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py similarity index 100% rename from task_sdk/tests/dags/super_basic_deferred_run.py rename to task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py diff --git a/task_sdk/tests/dags/super_basic_run.py b/task-sdk/tests/task_sdk/dags/super_basic_run.py similarity index 100% rename from task_sdk/tests/dags/super_basic_run.py rename to task-sdk/tests/task_sdk/dags/super_basic_run.py diff --git a/task_sdk/tests/api/__init__.py b/task-sdk/tests/task_sdk/definitions/__init__.py similarity index 100% rename from task_sdk/tests/api/__init__.py rename to task-sdk/tests/task_sdk/definitions/__init__.py diff --git a/task_sdk/tests/definitions/__init__.py b/task-sdk/tests/task_sdk/definitions/_internal/__init__.py similarity index 100% rename from task_sdk/tests/definitions/__init__.py rename to task-sdk/tests/task_sdk/definitions/_internal/__init__.py diff --git a/task_sdk/tests/definitions/_internal/test_templater.py b/task-sdk/tests/task_sdk/definitions/_internal/test_templater.py similarity index 100% rename from task_sdk/tests/definitions/_internal/test_templater.py rename to task-sdk/tests/task_sdk/definitions/_internal/test_templater.py diff --git a/task_sdk/tests/definitions/conftest.py b/task-sdk/tests/task_sdk/definitions/conftest.py similarity index 100% rename from task_sdk/tests/definitions/conftest.py rename to task-sdk/tests/task_sdk/definitions/conftest.py diff --git a/task_sdk/tests/definitions/test_asset.py b/task-sdk/tests/task_sdk/definitions/test_asset.py similarity index 100% rename from task_sdk/tests/definitions/test_asset.py rename to task-sdk/tests/task_sdk/definitions/test_asset.py diff --git a/task_sdk/tests/definitions/test_asset_decorators.py b/task-sdk/tests/task_sdk/definitions/test_asset_decorators.py similarity index 100% rename from task_sdk/tests/definitions/test_asset_decorators.py rename to task-sdk/tests/task_sdk/definitions/test_asset_decorators.py diff --git a/task_sdk/tests/definitions/test_baseoperator.py b/task-sdk/tests/task_sdk/definitions/test_baseoperator.py similarity index 100% rename from task_sdk/tests/definitions/test_baseoperator.py rename to task-sdk/tests/task_sdk/definitions/test_baseoperator.py diff --git a/task_sdk/tests/definitions/test_connections.py b/task-sdk/tests/task_sdk/definitions/test_connections.py similarity index 100% rename from task_sdk/tests/definitions/test_connections.py rename to task-sdk/tests/task_sdk/definitions/test_connections.py diff --git a/task_sdk/tests/definitions/test_context.py b/task-sdk/tests/task_sdk/definitions/test_context.py similarity index 100% rename from task_sdk/tests/definitions/test_context.py rename to task-sdk/tests/task_sdk/definitions/test_context.py diff --git a/task_sdk/tests/definitions/test_dag.py b/task-sdk/tests/task_sdk/definitions/test_dag.py similarity index 100% rename from task_sdk/tests/definitions/test_dag.py rename to task-sdk/tests/task_sdk/definitions/test_dag.py diff --git a/task_sdk/tests/definitions/test_macros.py b/task-sdk/tests/task_sdk/definitions/test_macros.py similarity index 100% rename from task_sdk/tests/definitions/test_macros.py rename to task-sdk/tests/task_sdk/definitions/test_macros.py diff --git a/task_sdk/tests/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py similarity index 100% rename from task_sdk/tests/definitions/test_mappedoperator.py rename to task-sdk/tests/task_sdk/definitions/test_mappedoperator.py diff --git a/task_sdk/tests/definitions/test_mixins.py b/task-sdk/tests/task_sdk/definitions/test_mixins.py similarity index 100% rename from task_sdk/tests/definitions/test_mixins.py rename to task-sdk/tests/task_sdk/definitions/test_mixins.py diff --git a/task_sdk/tests/definitions/test_param.py b/task-sdk/tests/task_sdk/definitions/test_param.py similarity index 100% rename from task_sdk/tests/definitions/test_param.py rename to task-sdk/tests/task_sdk/definitions/test_param.py diff --git a/task_sdk/tests/definitions/test_secrets_masker.py b/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py similarity index 100% rename from task_sdk/tests/definitions/test_secrets_masker.py rename to task-sdk/tests/task_sdk/definitions/test_secrets_masker.py diff --git a/task_sdk/tests/definitions/test_template.py b/task-sdk/tests/task_sdk/definitions/test_template.py similarity index 100% rename from task_sdk/tests/definitions/test_template.py rename to task-sdk/tests/task_sdk/definitions/test_template.py diff --git a/task_sdk/tests/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py similarity index 100% rename from task_sdk/tests/definitions/test_xcom_arg.py rename to task-sdk/tests/task_sdk/definitions/test_xcom_arg.py diff --git a/task_sdk/tests/definitions/_internal/__init__.py b/task-sdk/tests/task_sdk/execution_time/__init__.py similarity index 100% rename from task_sdk/tests/definitions/_internal/__init__.py rename to task-sdk/tests/task_sdk/execution_time/__init__.py diff --git a/task_sdk/tests/execution_time/conftest.py b/task-sdk/tests/task_sdk/execution_time/conftest.py similarity index 100% rename from task_sdk/tests/execution_time/conftest.py rename to task-sdk/tests/task_sdk/execution_time/conftest.py diff --git a/task_sdk/tests/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py similarity index 100% rename from task_sdk/tests/execution_time/test_context.py rename to task-sdk/tests/task_sdk/execution_time/test_context.py diff --git a/task_sdk/tests/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py similarity index 99% rename from task_sdk/tests/execution_time/test_supervisor.py rename to task-sdk/tests/task_sdk/execution_time/test_supervisor.py index ea639d9bdff71..66af26105fc7a 100644 --- a/task_sdk/tests/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -37,6 +37,7 @@ import pytest import tenacity from pytest_unordered import unordered +from task_sdk import FAKE_BUNDLE, make_client from uuid6 import uuid7 from airflow.executors.workloads import BundleInfo @@ -80,9 +81,6 @@ from airflow.sdk.execution_time.task_runner import CommsDecoder from airflow.utils import timezone, timezone as tz -from task_sdk.tests.api.test_client import make_client -from task_sdk.tests.execution_time.test_task_runner import FAKE_BUNDLE - if TYPE_CHECKING: import kgb diff --git a/task_sdk/tests/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py similarity index 99% rename from task_sdk/tests/execution_time/test_task_runner.py rename to task-sdk/tests/task_sdk/execution_time/test_task_runner.py index aee592122f8a8..dd223d69bfb0e 100644 --- a/task_sdk/tests/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -29,6 +29,7 @@ from unittest.mock import patch import pytest +from task_sdk import FAKE_BUNDLE from uuid6 import uuid7 from airflow.decorators import task as task_decorator @@ -98,8 +99,6 @@ if TYPE_CHECKING: from kgb import SpyAgency -FAKE_BUNDLE = BundleInfo(name="anything", version="any") - def get_inline_dag(dag_id: str, task: BaseOperator) -> DAG: """Creates an inline dag and returns it based on dag_id and task.""" diff --git a/task_sdk/tests/test_log.py b/task-sdk/tests/task_sdk/log/test_log.py similarity index 100% rename from task_sdk/tests/test_log.py rename to task-sdk/tests/task_sdk/log/test_log.py diff --git a/task_sdk/tests/execution_time/__init__.py b/task_sdk/tests/execution_time/__init__.py deleted file mode 100644 index 13a83393a9124..0000000000000 --- a/task_sdk/tests/execution_time/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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.