Task SDK Integration Testing Infrastructure

.dockerignore

@@ -46,6 +46,7 @@
 !docker-tests
 !helm-tests
 !kubernetes-tests
+!task-sdk-tests
 !shared/
 
 # Add scripts so that we can use them inside the container

.github/workflows/additional-prod-image-tests.yml

@@ -158,3 +158,33 @@ jobs:
         id: breeze
       - name: "Test docker-compose quick start"
         run: breeze testing docker-compose-tests
+
+  task-sdk-integration-tests:
+    timeout-minutes: 60
+    name: "Task SDK integration tests with PROD image"
+    runs-on: ${{ fromJSON(inputs.runners) }}
+    env:
+      PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}"
+      GITHUB_REPOSITORY: ${{ github.repository }}
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      GITHUB_USERNAME: ${{ github.actor }}
+      VERBOSE: "true"
+    steps:
+      - name: "Cleanup repo"
+        shell: bash
+        run: docker run -v "${GITHUB_WORKSPACE}:/workspace" -u 0:0 bash -c "rm -rf /workspace/*"
+      - name: "Checkout ${{ github.ref }} ( ${{ github.sha }} )"
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          fetch-depth: 2
+          persist-credentials: false
+      - name: "Prepare breeze & PROD image: ${{ env.PYTHON_MAJOR_MINOR_VERSION }}"
+        uses: ./.github/actions/prepare_breeze_and_image
+        with:
+          platform: ${{ inputs.platform }}
+          image-type: "prod"
+          python: ${{ env.PYTHON_MAJOR_MINOR_VERSION }}
+          use-uv: ${{ inputs.use-uv }}
+        id: breeze
+      - name: "Run Task SDK integration tests"
+        run: breeze testing task-sdk-integration-tests

Dockerfile

@@ -904,6 +904,7 @@ function install_from_sources() {
         installation_command_flags=" --editable .[${AIRFLOW_EXTRAS}] \
               --editable ./airflow-core --editable ./task-sdk --editable ./airflow-ctl \
               --editable ./kubernetes-tests --editable ./docker-tests --editable ./helm-tests \
+              --editable ./task-sdk-tests \
               --editable ./devel-common[all] --editable ./dev \
               --group dev --group docs --group docs-gen --group leveldb"
         local -a projects_with_devel_dependencies

Dockerfile.ci

@@ -687,6 +687,7 @@ function install_from_sources() {
         installation_command_flags=" --editable .[${AIRFLOW_EXTRAS}] \
               --editable ./airflow-core --editable ./task-sdk --editable ./airflow-ctl \
               --editable ./kubernetes-tests --editable ./docker-tests --editable ./helm-tests \
+              --editable ./task-sdk-tests \
               --editable ./devel-common[all] --editable ./dev \
               --group dev --group docs --group docs-gen --group leveldb"
         local -a projects_with_devel_dependencies

contributing-docs/09_testing.rst

@@ -44,6 +44,9 @@ includes:
 * `System tests <testing/system_tests.rst>`__ are automatic tests that use external systems like
   Google Cloud and AWS. These tests are intended for an end-to-end DAG execution.
 
+* `Task SDK integration tests <testing/task_sdk_integration_tests.rst>`__ are specialized tests that verify
+  the integration between the Apache Airflow Task SDK package and a running Airflow instance.
+
 You can also run other kinds of tests when you are developing Airflow packages:
 
 * `Testing distributions <testing/testing_distributions.rst>`__ is a document that describes how to

contributing-docs/testing/task_sdk_integration_tests.rst

@@ -0,0 +1,215 @@
+ .. 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.
+
+Task SDK Integration Tests
+==========================
+
+Task SDK Integration Tests are a specialized type of test that verify the integration between the
+`Apache Airflow Task SDK <../task-sdk/>`__ and a running Airflow instance (mainly the execution API server).
+These tests ensure that the Task SDK can properly communicate with Airflow's execution API server
+and that the integration between the two works correctly in a realistic environment.
+
+What are Task SDK Integration Tests?
+------------------------------------
+
+The Task SDK Integration Tests differ from regular unit tests in several key ways:
+
+**Purpose & Scope**
+  These tests verify that the Task SDK package can successfully communicate with a complete Airflow
+  environment, mainly the Airflow Execution API server. They test the actual
+  integration points between the SDK and Airflow core, which is also known as Task Execution Interface.
+
+**Real Environment Testing**
+  Unlike unit tests that use mocks and stubs, these integration tests spin up actual Airflow services
+  using Docker Compose to create a realistic testing environment that closely mirrors real deployments.
+
+**Cross-Package Validation**
+  Since the Task SDK is distributed as a separate Python package (``apache-airflow-task-sdk``),
+  these tests ensure that the SDK works correctly with different versions of Airflow core and
+  that API compatibility is maintained.
+
+**End-to-End Communication**
+  The tests verify complete request/response cycles between the Task SDK API client and Airflow's
+  execution API, including authentication, API versioning, and error handling (some of it is yet to come).
+
+Why Task SDK Integration?
+-------------------------
+
+**API Compatibility Assurance**
+  The Task SDK communicates with Airflow through well defined APIs for task execution. These tests ensure that
+  changes to either the SDK or Airflow core don't break the interface contract between the two.
+
+**Real World Scenario Testing**
+  While unit tests verify individual components work correctly, integration tests validate
+  that the entire system works together as expected in deployment scenarios.
+
+**Quicker Interface Issue resolution**
+  These tests catch integration issues early in the development cycle, preventing breaking
+  changes reaching a release.
+
+**Version Compatibility Matrix**
+  As both Airflow and the Task SDK evolve, these tests help ensure compatibility across
+  different version combinations (to come soon).
+
+Running Task SDK Integration Tests
+----------------------------------
+
+There are multiple ways to run Task SDK Integration Tests depending based on your preferences.
+
+Using Breeze
+............
+
+The simplest way to run Task SDK Integration Tests is using Breeze, which provides CI like
+reproducibility:
+
+.. code-block:: bash
+
+   # Run all Task SDK integration tests
+   breeze testing task-sdk-integration-tests
+
+   # Run specific test files
+   breeze testing task-sdk-integration-tests task_sdk_tests/test_task_sdk_health.py
+
+   # Run with custom Docker image
+   DOCKER_IMAGE=my-custom-airflow-imahge:latest breeze testing task-sdk-integration-tests
+
+Running in Your Current Virtual Environment
+...........................................
+
+Since you're already working in the Airflow repository, you can run Task SDK Integration Tests
+directly:
+
+**Run Tests**
+
+.. code-block:: bash
+
+   # Navigate to task-sdk-tests directory and run tests
+   cd task-sdk-tests/
+   uv run pytest -s
+
+   # Run specific test file
+   cd task-sdk-tests/
+   uv run pytest tests/task_sdk_tests/test_task_sdk_health.py -s
+
+   # Keep containers running for debugging
+   cd task-sdk-tests/
+   SKIP_DOCKER_COMPOSE_DELETION=1 uv run pytest -s
+
+**Optional: Set Custom Docker Image**
+
+.. code-block:: bash
+
+   # Use a different Airflow image for testing
+   cd task-sdk-tests/
+   DOCKER_IMAGE=my-custom-airflow:latest uv run pytest -s
+
+
+Debugging Failed Tests
+......................
+
+When tests fail, the logs from all running containers are automatically dumped to the console
+and the Docker Compose deployment is shut down. To debug issues more effectively:
+
+.. code-block:: bash
+
+   # Run with maximum verbosity
+   cd task-sdk-tests/
+   uv run pytest tests/task_sdk_tests/ -vvv -s --tb=long
+
+   # Keep containers running for inspection (local environment)
+   cd task-sdk-tests/
+   SKIP_DOCKER_COMPOSE_DELETION=1 uv run pytest tests/task_sdk_tests/test_task_sdk_health.py::test_task_sdk_health
+
+   # Keep containers running for inspection (using Breeze)
+   breeze testing task-sdk-integration-tests --skip-docker-compose-deletion
+
+   # Inspect container logs (when containers are still running)
+   cd task-sdk-tests/docker
+   docker-compose logs airflow-apiserver
+   docker-compose logs airflow-scheduler
+   docker-compose logs postgres
+
+   # Access running containers for interactive debugging
+   docker-compose exec airflow-apiserver bash
+
+.. tip::
+   **Container Cleanup Control**: By default, the Docker Compose deployment is deleted after tests
+   complete to keep your system clean. To keep containers running for debugging:
+
+   - **Local environment**: Export ``SKIP_DOCKER_COMPOSE_DELETION=1`` before running tests
+   - **Breeze environment**: Use the ``--skip-docker-compose-deletion`` flag
+
+   Remember to manually clean up containers when done: ``cd task-sdk-tests/docker && docker-compose down -v``
+
+Testing Custom Airflow Images
+..............................
+
+To test the Task SDK against custom Airflow builds:
+
+.. code-block:: bash
+
+   # Build your custom Airflow image first
+   cd /path/to/airflow
+   docker build -t my-custom-airflow:latest -f Dockerfile .
+
+   # Use custom image for integration tests
+   export DOCKER_IMAGE=my-custom-airflow:latest
+   cd task-sdk-tests
+   uv run pytest tests/task_sdk_tests/
+
+Common Issues and Solutions
+---------------------------
+
+**Port Conflicts**
+  If port 8080 is already in use, change the host port:
+
+  .. code-block:: bash
+
+     export HOST_PORT=localhost:9090
+     export TASK_SDK_HOST_PORT=localhost:9090
+
+**Memory Issues**
+  These tests require sufficient memory for multiple containers. Ensure Docker has at least 4GB RAM allocated.
+
+Files and Directories
+---------------------
+
+The Task SDK Integration Tests are organized as follows:
+
+.. code-block::
+
+   task-sdk-tests/
+   ├── pyproject.toml                    # Test package configuration and dependencies
+   ├── docker/
+   │   └── docker-compose.yaml           # Airflow services configuration
+   └── tests/
+       └── task_sdk_tests/
+           ├── conftest.py                # Test configuration and setup
+           ├── constants.py               # Test constants and configuration
+           ├── test_task_sdk_health.py    # Main integration test
+           └── __init__.py
+
+**Key Files:**
+
+- **docker-compose.yaml**: Defines the complete Airflow environment (postgres, scheduler, api-server)
+- **test_task_sdk_health.py**: Main test that verifies Task SDK can communicate with Airflow API
+- **conftest.py**: Handles Task SDK installation and test environment setup
+- **constants.py**: Configuration constants for Docker images, ports, and API versions
+
+-----
+
+For other kinds of tests look at `Testing document <../09_testing.rst>`__

dev/breeze/doc/05_test_commands.rst

@@ -317,6 +317,25 @@ through ``breeze testing docker-compose-tests`` command.
 
 The docker-compose tests are in ``docker-tests/`` folder in the main repo.
 
+Running task-sdk integration tests
+..................................
+
+You can use Breeze to run the task sdk integration tests. Those tests are run using Production image by default
+and the tests are running with the docker-compose we have for task-sdk tests.
+
+.. image:: ./images/output_testing_docker-compose-tests.svg
+  :target: https://raw.githubusercontent.com/apache/airflow/main/dev/breeze/images/output_testing_task-sdk-integration-tests.svg
+  :width: 100%
+  :alt: Breeze testing task-sdk-integration-tests
+
+You can also iterate over those tests with pytest command but unlike regular unit tests, they need to be run in
+a local venv. You can build the prod image with breeze and that will be used by default if present to run the tests.
+
+You can override the ``DOCKER_IMAGE`` environment variable to point to the image to test using the
+``breeze testing task-sdk-integration-tests`` command.
+
+The task-sdk tests are in ``task-sdk-tests/`` folder in the main repo.
+
 Running Kubernetes tests
 ------------------------