Skip to content

Add task SDK integration test foundation with JWT auth and session fixtures#56139

Merged
amoghrajesh merged 3 commits into
apache:mainfrom
astronomer:infra-task-sdk-testing
Oct 14, 2025
Merged

Add task SDK integration test foundation with JWT auth and session fixtures#56139
amoghrajesh merged 3 commits into
apache:mainfrom
astronomer:infra-task-sdk-testing

Conversation

@amoghrajesh

@amoghrajesh amoghrajesh commented Sep 26, 2025

Copy link
Copy Markdown
Contributor

Purpose

We introduced a test integration framework for task SDK through: #53332. This PR establishes the foundational infrastructure for comprehensive integration testing so that the current test suite can be expanded further.

Through this, I aim to creates reusable components and patterns to enable will enable systematic validation of all potentially all execution API endpoints used by the task SDK rest API client.

What's being done?

Core infrastructure

  1. JWT Authentication Plugin (jwt_plugin.py) - Generates valid JWT tokens for Execution API authentication using task instance IDs, with hourly expiry. We need to do this because we will be calling execution API endpoints which require a token for task instance and using this token we can call API endpoints on this task
  2. Session-Scoped Fixtures (conftest.py) - Efficient test setup that runs docker-compose once per session, triggers test DAG, and provides authenticated clients as needed.

Docker compose improvements

  1. DAG Processor Dependency - Updated docker-compose to ensure proper component startup order (dag-processor → scheduler). Earlier, the dag processor didn't exist, created a minimal dag that runs for long enough and will be useful for testing situations.
  2. Environment Configuration - Proper JWT secret and API URL configuration for test environment

Test Refactoring

  1. Health Test Updated - Refactored existing health check to use session fixtures instead of per-test docker-compose

Task SDK integration testing flow now

Session Start (Once):

  1. Install Task SDK
  2. Start docker-compose (postgres → init → dag-processor → scheduler → api-server)
  3. Get auth token from SimpleAuthManager
  4. Auto-unpause test_dag
  5. Trigger test_dag via Core API
  6. Wait for task instances
  7. Generate JWT token for task instance
  8. Create authenticated SDK client

Individual Tests (Reuse session):

  • Use pre-configured authenticated client
  • Test specific Execution API endpoints
  • Validate responses

Session End:

  • Clean up docker-compose

What's Next

This foundation enables adding tests for all Execution API endpoints. My idea is that I will create a structure that can be extended easily to add some tests based on expecatations defined in a future issue so that ANYONE can contribute to it as needed.

Community Contributions:
The standardized patterns make it easy for community members to contribute endpoint-specific tests following established templates.


^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

Comment thread task-sdk-tests/tests/task_sdk_tests/constants.py Outdated
Comment thread task-sdk-tests/tests/task_sdk_tests/constants.py
@amoghrajesh amoghrajesh requested a review from shahar1 October 14, 2025 12:56
@amoghrajesh amoghrajesh merged commit ac14b47 into apache:main Oct 14, 2025
51 checks passed
@amoghrajesh amoghrajesh deleted the infra-task-sdk-testing branch October 14, 2025 16:26
TyrellHaywood pushed a commit to TyrellHaywood/airflow that referenced this pull request Oct 22, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

No open projects

Development

Successfully merging this pull request may close these issues.

3 participants