You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "amoghrajesh (via GitHub)" <gi...@apache.org> on 2025/11/07 07:48:08 UTC
[PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
amoghrajesh opened a new pull request, #58027:
URL: https://github.com/apache/airflow/pull/58027
<!--
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.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
Adding tests for asset operations by virtue of doing the following:
1. Adding similar test setup for asset dag trigger etc as the normal dag operations added in #56139
2. Adding an example test dag for assets that will be triggered manually through the tests to test for creation of assets
3. Covers the `get` call for now, both positive and negative cases, will add stubs for the rest in a follow up
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
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](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #58027:
URL: https://github.com/apache/airflow/pull/58027#discussion_r2506089209
##########
task-sdk-tests/tests/task_sdk_tests/conftest.py:
##########
@@ -182,9 +183,55 @@ def pytest_sessionstart(session):
raise
-@pytest.fixture(scope="session")
-def airflow_test_setup(docker_compose_setup):
- """Fixed session-scoped fixture that matches UI behavior."""
+def setup_dag_and_get_client(
+ *,
+ dag_id: str,
+ headers: dict[str, str],
+ auth_token: str | None = None,
+ task_id_filter: str | Callable[[dict], bool] | None = None,
+ wait_for_task_state: str | None = None,
+ wait_for_dag_state: str | None = None,
+ wait_timeout: int = 60,
+ additional_metadata: dict | None = None,
+) -> dict:
+ """
+ Utility to set up a DAG run and create an SDK client.
+
+ This function handles the common pattern of:
+ 1. Getting DAG status
+ 2. Unpausing DAG (if needed)
+ 3. Triggering a DAG run
+ 4. Waiting for task instances or DAG run state
+ 5. Acquiring the task instance ID for the triggered DAG run
+ 6. Generating JWT token for that task instance
+ 7. Creating a task SDK client with that JWT token
+
+ Args:
Review Comment:
Small nits. If I remember correctly, we are using `:param` style doc string.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #58027:
URL: https://github.com/apache/airflow/pull/58027#discussion_r2504127709
##########
task-sdk-tests/tests/task_sdk_tests/conftest.py:
##########
@@ -320,3 +333,111 @@ def sdk_client(airflow_test_setup):
def core_api_headers(airflow_test_setup):
"""Get Core API headers from setup."""
return airflow_test_setup["core_api_headers"]
+
+
+@pytest.fixture(scope="session")
+def sdk_client_for_assets(asset_test_setup):
+ """Get SDK client for asset tests (doesn't require test_dag)."""
+ return asset_test_setup["sdk_client"]
+
+
+@pytest.fixture(scope="session")
+def asset_test_setup(docker_compose_setup, airflow_ready):
+ """Setup assets for testing by triggering asset_producer_dag."""
+ import time
+
+ import requests
+
+ from airflow.sdk.api.client import Client
+ from airflow.sdk.timezone import utcnow
+ from task_sdk_tests.jwt_plugin import generate_jwt_token
+
+ headers = airflow_ready["headers"]
+
+ # For asset test setup, we need to trigger asset_producer_dag to create test asset
+ console.print("[yellow]Checking asset_producer_dag status...")
+ dag_response = requests.get("http://localhost:8080/api/v2/dags/asset_producer_dag", headers=headers)
+ dag_response.raise_for_status()
+ dag_data = dag_response.json()
+
+ if dag_data.get("is_paused", True):
+ console.print("[yellow]Unpausing asset_producer_dag...")
+ unpause_response = requests.patch(
+ "http://localhost:8080/api/v2/dags/asset_producer_dag", json={"is_paused": False}, headers=headers
+ )
+ unpause_response.raise_for_status()
+ console.print("[green]asset_producer_dag unpaused")
+
+ logical_date = utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ")[:-3] + "Z"
+ payload = {"conf": {}, "logical_date": logical_date}
+
+ console.print("[yellow]Triggering asset_producer_dag...")
+ trigger_response = requests.post(
+ "http://localhost:8080/api/v2/dags/asset_producer_dag/dagRuns",
+ json=payload,
+ headers=headers,
+ timeout=30,
+ )
+
+ trigger_response.raise_for_status()
+ dag_run_data = trigger_response.json()
+ dag_run_id = dag_run_data["dag_run_id"]
+
+ console.print(f"[green]asset_producer_dag triggered: {dag_run_id}")
+
+ console.print("[yellow]Waiting for asset_producer_dag to complete...")
+ final_state = None
+ for attempt in range(60):
+ try:
+ dr_response = requests.get(
+ f"http://localhost:8080/api/v2/dags/asset_producer_dag/dagRuns/{dag_run_id}", headers=headers
+ )
+ dr_response.raise_for_status()
+ dr_data = dr_response.json()
+ state = dr_data.get("state")
+
+ if state == "success":
+ console.print("[green]asset_producer_dag completed successfully!")
+ final_state = state
+ break
+ if state in ["failed", "skipped"]:
+ raise RuntimeError(f"asset_producer_dag ended in state: {state}")
+ console.print(
+ f"[blue]Waiting for asset_producer_dag to complete (attempt {attempt + 1}/60, state: {state})"
+ )
+
+ except Exception as e:
+ console.print(f"[yellow]DAG run check failed: {e}")
+
+ time.sleep(2)
+
+ if final_state != "success":
+ raise TimeoutError("asset_producer_dag did not complete successfully within timeout period")
+
+ # Get task instance for testing - wait for long_running_task to be RUNNING
+ console.print("[yellow]Getting task instance ID from asset_producer_dag...")
+ ti_url = f"http://localhost:8080/api/v2/dags/asset_producer_dag/dagRuns/{dag_run_id}/taskInstances"
+ ti_response = requests.get(ti_url, headers=headers, timeout=10)
+ ti_response.raise_for_status()
+
+ task_instances = ti_response.json().get("task_instances", [])
+ ti_id = None
+ for ti in task_instances:
+ if ti.get("task_id") == "produce_asset":
+ ti_id = ti.get("id")
+ break
+
+ if not ti_id:
+ raise RuntimeError("Could not find task instance ID from asset_producer_dag")
+
+ console.print(f"[green]Found task instance ID: {ti_id}")
+
+ jwt_token = generate_jwt_token(ti_id)
+ sdk_client = Client(base_url=f"http://{TASK_SDK_HOST_PORT}/execution", token=jwt_token)
+
Review Comment:
Non-blocking nit.
Would it be better to have a common utils for dag setup?
Since `airflow_test_setup` and `asset_test_setup` will both
1. get dag
2. unpause dag
3. trigger dag
4. wait for some TIs show up
5. generate TI specific JWT Token
6. return client
##########
task-sdk-tests/tests/task_sdk_tests/test_asset_operations.py:
##########
@@ -0,0 +1,65 @@
+#
+# 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.
+"""
+Integration tests for Asset operations.
+
+These tests validate the Execution API endpoints for Asset operations:
+- get(): Get asset by name
+"""
+
+from __future__ import annotations
+
+from airflow.sdk.api.datamodels._generated import AssetResponse
+from airflow.sdk.execution_time.comms import ErrorResponse
+from task_sdk_tests import console
+
Review Comment:
Should we validate whether `asset_consumer_dag` dag is triggered or not?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #58027:
URL: https://github.com/apache/airflow/pull/58027#discussion_r2504815052
##########
task-sdk-tests/tests/task_sdk_tests/conftest.py:
##########
@@ -320,3 +333,111 @@ def sdk_client(airflow_test_setup):
def core_api_headers(airflow_test_setup):
"""Get Core API headers from setup."""
return airflow_test_setup["core_api_headers"]
+
+
+@pytest.fixture(scope="session")
+def sdk_client_for_assets(asset_test_setup):
+ """Get SDK client for asset tests (doesn't require test_dag)."""
+ return asset_test_setup["sdk_client"]
+
+
+@pytest.fixture(scope="session")
+def asset_test_setup(docker_compose_setup, airflow_ready):
+ """Setup assets for testing by triggering asset_producer_dag."""
+ import time
+
+ import requests
+
+ from airflow.sdk.api.client import Client
+ from airflow.sdk.timezone import utcnow
+ from task_sdk_tests.jwt_plugin import generate_jwt_token
+
+ headers = airflow_ready["headers"]
+
+ # For asset test setup, we need to trigger asset_producer_dag to create test asset
+ console.print("[yellow]Checking asset_producer_dag status...")
+ dag_response = requests.get("http://localhost:8080/api/v2/dags/asset_producer_dag", headers=headers)
+ dag_response.raise_for_status()
+ dag_data = dag_response.json()
+
+ if dag_data.get("is_paused", True):
+ console.print("[yellow]Unpausing asset_producer_dag...")
+ unpause_response = requests.patch(
+ "http://localhost:8080/api/v2/dags/asset_producer_dag", json={"is_paused": False}, headers=headers
+ )
+ unpause_response.raise_for_status()
+ console.print("[green]asset_producer_dag unpaused")
+
+ logical_date = utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ")[:-3] + "Z"
+ payload = {"conf": {}, "logical_date": logical_date}
+
+ console.print("[yellow]Triggering asset_producer_dag...")
+ trigger_response = requests.post(
+ "http://localhost:8080/api/v2/dags/asset_producer_dag/dagRuns",
+ json=payload,
+ headers=headers,
+ timeout=30,
+ )
+
+ trigger_response.raise_for_status()
+ dag_run_data = trigger_response.json()
+ dag_run_id = dag_run_data["dag_run_id"]
+
+ console.print(f"[green]asset_producer_dag triggered: {dag_run_id}")
+
+ console.print("[yellow]Waiting for asset_producer_dag to complete...")
+ final_state = None
+ for attempt in range(60):
+ try:
+ dr_response = requests.get(
+ f"http://localhost:8080/api/v2/dags/asset_producer_dag/dagRuns/{dag_run_id}", headers=headers
+ )
+ dr_response.raise_for_status()
+ dr_data = dr_response.json()
+ state = dr_data.get("state")
+
+ if state == "success":
+ console.print("[green]asset_producer_dag completed successfully!")
+ final_state = state
+ break
+ if state in ["failed", "skipped"]:
+ raise RuntimeError(f"asset_producer_dag ended in state: {state}")
+ console.print(
+ f"[blue]Waiting for asset_producer_dag to complete (attempt {attempt + 1}/60, state: {state})"
+ )
+
+ except Exception as e:
+ console.print(f"[yellow]DAG run check failed: {e}")
+
+ time.sleep(2)
+
+ if final_state != "success":
+ raise TimeoutError("asset_producer_dag did not complete successfully within timeout period")
+
+ # Get task instance for testing - wait for long_running_task to be RUNNING
+ console.print("[yellow]Getting task instance ID from asset_producer_dag...")
+ ti_url = f"http://localhost:8080/api/v2/dags/asset_producer_dag/dagRuns/{dag_run_id}/taskInstances"
+ ti_response = requests.get(ti_url, headers=headers, timeout=10)
+ ti_response.raise_for_status()
+
+ task_instances = ti_response.json().get("task_instances", [])
+ ti_id = None
+ for ti in task_instances:
+ if ti.get("task_id") == "produce_asset":
+ ti_id = ti.get("id")
+ break
+
+ if not ti_id:
+ raise RuntimeError("Could not find task instance ID from asset_producer_dag")
+
+ console.print(f"[green]Found task instance ID: {ti_id}")
+
+ jwt_token = generate_jwt_token(ti_id)
+ sdk_client = Client(base_url=f"http://{TASK_SDK_HOST_PORT}/execution", token=jwt_token)
+
Review Comment:
Sure, not a bad idea at all.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #58027:
URL: https://github.com/apache/airflow/pull/58027#discussion_r2504818615
##########
task-sdk-tests/tests/task_sdk_tests/test_asset_operations.py:
##########
@@ -0,0 +1,65 @@
+#
+# 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.
+"""
+Integration tests for Asset operations.
+
+These tests validate the Execution API endpoints for Asset operations:
+- get(): Get asset by name
+"""
+
+from __future__ import annotations
+
+from airflow.sdk.api.datamodels._generated import AssetResponse
+from airflow.sdk.execution_time.comms import ErrorResponse
+from task_sdk_tests import console
+
Review Comment:
Good shout, actually we do not even need the consumer_dag because the tests are testing for "asset operations" -> does the asset exist, fetching with name or uri, etc and later the asset event which will be emitted irrespective of consumer being there or not, let me delete that dag entirely.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #58027:
URL: https://github.com/apache/airflow/pull/58027#discussion_r2506084621
##########
task-sdk-tests/tests/task_sdk_tests/test_asset_operations.py:
##########
@@ -0,0 +1,65 @@
+#
+# 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.
+"""
+Integration tests for Asset operations.
+
+These tests validate the Execution API endpoints for Asset operations:
+- get(): Get asset by name
+"""
+
+from __future__ import annotations
+
+from airflow.sdk.api.datamodels._generated import AssetResponse
+from airflow.sdk.execution_time.comms import ErrorResponse
+from task_sdk_tests import console
+
Review Comment:
Gotcha! Thanks for the clarification.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #58027:
URL: https://github.com/apache/airflow/pull/58027#discussion_r2506918050
##########
task-sdk-tests/tests/task_sdk_tests/conftest.py:
##########
@@ -182,9 +183,55 @@ def pytest_sessionstart(session):
raise
-@pytest.fixture(scope="session")
-def airflow_test_setup(docker_compose_setup):
- """Fixed session-scoped fixture that matches UI behavior."""
+def setup_dag_and_get_client(
+ *,
+ dag_id: str,
+ headers: dict[str, str],
+ auth_token: str | None = None,
+ task_id_filter: str | Callable[[dict], bool] | None = None,
+ wait_for_task_state: str | None = None,
+ wait_for_dag_state: str | None = None,
+ wait_timeout: int = 60,
+ additional_metadata: dict | None = None,
+) -> dict:
+ """
+ Utility to set up a DAG run and create an SDK client.
+
+ This function handles the common pattern of:
+ 1. Getting DAG status
+ 2. Unpausing DAG (if needed)
+ 3. Triggering a DAG run
+ 4. Waiting for task instances or DAG run state
+ 5. Acquiring the task instance ID for the triggered DAG run
+ 6. Generating JWT token for that task instance
+ 7. Creating a task SDK client with that JWT token
+
+ Args:
Review Comment:
Thats a good point but not to stress over that because we are not gonna ever generate docs for this distro at all. That would be useful if we did, this is more of an information sake docstring.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add test setup and example asset operation tests to task SDK integration tests [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh merged PR #58027:
URL: https://github.com/apache/airflow/pull/58027
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org