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
[email protected](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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
