This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f760ac1c03b Stop airflowctl e2e xcom tests failing on a slow compose
stack (#72260)
f760ac1c03b is described below
commit f760ac1c03b7277f3cb69071b8ead56d27ecbff7
Author: rjgoyln <[email protected]>
AuthorDate: Sun Aug 30 23:09:11 2026 +0800
Stop airflowctl e2e xcom tests failing on a slow compose stack (#72260)
The xcom commands can only run once their Dag run has finished, and the
wait that enforces that gave a freshly booted PROD compose stack only 60
seconds. Every sample cost a full airflowctl subprocess, so a slow machine
got barely a dozen looks in before the budget ran out. Because all five
xcom commands share one Dag run, a single slow boot failed all of them at
once, behind a message that said nothing about how far the run had got.
---
.../tests/airflowctl_tests/conftest.py | 157 ++++++++++-----
.../tests/airflowctl_tests/constants.py | 11 +-
.../tests/airflowctl_tests/test_dag_run_wait.py | 222 +++++++++++++++++++++
contributing-docs/testing/airflow_ctl_tests.rst | 12 ++
4 files changed, 350 insertions(+), 52 deletions(-)
diff --git a/airflow-ctl-tests/tests/airflowctl_tests/conftest.py
b/airflow-ctl-tests/tests/airflowctl_tests/conftest.py
index 617bd539df9..227f9d7d431 100644
--- a/airflow-ctl-tests/tests/airflowctl_tests/conftest.py
+++ b/airflow-ctl-tests/tests/airflowctl_tests/conftest.py
@@ -16,12 +16,12 @@
# under the License.
from __future__ import annotations
-import json
import os
import re
import subprocess
import sys
import time
+from collections.abc import Callable, Iterator
from subprocess import PIPE, STDOUT, Popen
import pytest
@@ -31,12 +31,18 @@ from python_on_whales import DockerClient, docker
from airflowctl_tests import console
from airflowctl_tests.constants import (
AIRFLOW_ROOT_PATH,
+ API_PASSWORD,
+ API_USERNAME,
+ DAG_RUN_WAIT_TIMEOUT,
+ DAG_RUN_WAIT_TIMEOUT_ENV,
DOCKER_COMPOSE_FILE_PATH,
+ DOCKER_COMPOSE_HOST_PORT,
DOCKER_IMAGE,
LOGIN_COMMAND,
LOGIN_OUTPUT,
)
+from tests_common.test_utils.api_client_helpers import generate_access_token
from tests_common.test_utils.fernet import generate_fernet_key_string
# XCom add/edit/delete race against task execution: when the target task
transitions to
@@ -49,73 +55,126 @@ from tests_common.test_utils.fernet import
generate_fernet_key_string
# XComs survive the rest of the xcom commands.
_XCOM_TARGET_PATTERN =
re.compile(r'^xcom\s+(?:add|get|list|edit|delete)\s+(\S+)\s+"(manual__[^"]+)"')
_DAG_RUN_TERMINAL_STATES = frozenset({"success", "failed"})
+_API_REQUEST_TIMEOUT = 30
+_POLL_INITIAL_DELAY = 0.5
+_POLL_MAX_DELAY = 10.0
-def _airflowctl_dag_run_state(dag_id: str, dag_run_id: str, env_vars: dict,
skip_login: bool) -> str | None:
- """Return the current state of a Dag run via airflowctl, or None if
unparsable."""
- host_envs = os.environ.copy()
- host_envs.update(env_vars)
+class _CtlTestState:
+ docker_client: DockerClient | None = None
+ access_token: str | None = None
+
+
+def _get_access_token() -> str:
+ """Log in once per pytest session; every REST call below reuses the
token."""
+ if _CtlTestState.access_token is None:
+ _CtlTestState.access_token = generate_access_token(
+ API_USERNAME, API_PASSWORD, DOCKER_COMPOSE_HOST_PORT
+ )
+ return _CtlTestState.access_token
+
+
+def _request_api(path: str) -> dict:
+ """GET a public API path on the compose stack, e.g.
``dags/my_dag/dagRuns/run_id``."""
+ response = requests.get(
+ f"http://{DOCKER_COMPOSE_HOST_PORT}/api/v2/{path}";,
+ headers={"Authorization": f"Bearer {_get_access_token()}"},
+ timeout=_API_REQUEST_TIMEOUT,
+ )
+ response.raise_for_status()
+ return response.json()
- get_cmd = f'airflowctl dagrun get {dag_id} "{dag_run_id}" -o json'
- if not skip_login:
- get_cmd = f"airflowctl {LOGIN_COMMAND} && {get_cmd}"
- proc = Popen(get_cmd.encode(), stdout=PIPE, stderr=STDOUT, shell=True,
env=host_envs)
+def _find_dag_run_state(dag_id: str, dag_run_id: str) -> str | None:
+ """Return the current state of a Dag run, or None while the API cannot
answer."""
try:
- out, _ = proc.communicate(timeout=20)
- except subprocess.TimeoutExpired:
- proc.kill()
- return None
- out_str = out.decode()
- if LOGIN_OUTPUT in out_str:
- out_str = out_str.split(f"{LOGIN_OUTPUT}\n", 1)[-1].strip()
- start, end = out_str.find("{"), out_str.rfind("}")
- if start == -1 or end == -1:
+ return _request_api(f"dags/{dag_id}/dagRuns/{dag_run_id}").get("state")
+ except requests.exceptions.HTTPError as exc:
+ # A rejected request means a wrong run id or wrong credentials;
polling on until
+ # the budget runs out would bury that behind a timeout.
+ if exc.response is not None and exc.response.status_code >= 500:
+ return None
+ raise
+ except requests.exceptions.RequestException:
return None
+
+
+def _describe_task_instances(dag_id: str, dag_run_id: str) -> str:
+ """Render the Dag run's ``task_id=state`` pairs for timeout diagnostics."""
try:
- return json.loads(out_str[start : end + 1]).get("state")
- except json.JSONDecodeError:
- return None
+ task_instances =
_request_api(f"dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances")["task_instances"]
+ except (requests.exceptions.RequestException, KeyError) as exc:
+ return f"unavailable ({exc})"
+ return ", ".join(f"{ti['task_id']}={ti['state']}" for ti in
task_instances) or "none"
+
+
+def _compute_poll_delays() -> Iterator[float]:
+ """Sample often while the run is likely to finish soon, then ease off."""
+ delay = _POLL_INITIAL_DELAY
+ while True:
+ yield delay
+ delay = min(delay * 2, _POLL_MAX_DELAY)
def _wait_for_dag_run_terminal_state(
dag_id: str,
dag_run_id: str,
- env_vars: dict,
- skip_login: bool,
- timeout: int = 60,
+ timeout: float = DAG_RUN_WAIT_TIMEOUT,
) -> None:
"""Block until the Dag run reaches success/failed, or raise
TimeoutError."""
deadline = time.monotonic() + timeout
- last_state: str | None = None
- while time.monotonic() < deadline:
- last_state = _airflowctl_dag_run_state(dag_id, dag_run_id, env_vars,
skip_login)
- if last_state in _DAG_RUN_TERMINAL_STATES:
+ delays = _compute_poll_delays()
+ while True:
+ state = _find_dag_run_state(dag_id, dag_run_id)
+ if state in _DAG_RUN_TERMINAL_STATES:
return
- time.sleep(1)
- raise TimeoutError(
- f"Dag run {dag_id}/{dag_run_id} did not reach terminal state in
{timeout}s "
- f"(last seen state: {last_state})"
- )
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ raise TimeoutError(
+ f"Dag run {dag_id}/{dag_run_id} did not reach a terminal state
in {timeout}s "
+ f"(Dag run state: {state}; task instances:
{_describe_task_instances(dag_id, dag_run_id)}). "
+ f"Set {DAG_RUN_WAIT_TIMEOUT_ENV} to give a slow machine a
longer budget."
+ )
+ time.sleep(min(next(delays), remaining))
+
+
+def _build_dag_run_waiter() -> Callable[[str, str], None]:
+ """Spend the timeout budget once per Dag run and replay that outcome to
later commands."""
+ outcomes: dict[tuple[str, str], str | None] = {}
+
+ def _wait(dag_id: str, dag_run_id: str) -> None:
+ key = (dag_id, dag_run_id)
+ if key not in outcomes:
+ try:
+ _wait_for_dag_run_terminal_state(dag_id, dag_run_id)
+ except TimeoutError as exc:
+ outcomes[key] = str(exc)
+ raise
+ outcomes[key] = None
+ elif outcomes[key] is not None:
+ # The budget is spent, but one cheap look keeps a run that
finished late from
+ # skipping the rest of the xcom commands.
+ if _find_dag_run_state(dag_id, dag_run_id) not in
_DAG_RUN_TERMINAL_STATES:
+ pytest.skip(f"Dag run {dag_id}/{dag_run_id} is unusable:
{outcomes[key]}")
+ outcomes[key] = None
+
+ return _wait
[email protected](scope="module")
[email protected](scope="session")
+def wait_for_dag_run():
+ """Fixture that provides a helper to wait for a Dag run to become
terminal."""
+ return _build_dag_run_waiter()
+
+
[email protected](scope="session")
def api_token():
- url = "http://localhost:8080/auth/token";
- payload = {"username": "airflow", "password": "airflow"}
- try:
- response = requests.post(url, json=payload)
- response.raise_for_status()
- token = response.json().get("access_token")
- if not token:
- raise ValueError("Response did not contain an access_token")
- return token
- except requests.exceptions.RequestException as e:
- pytest.fail(f"Failed to obtain token: {e}")
+ """Fixture that provides the API token shared by the whole test session."""
+ return _get_access_token()
@pytest.fixture
-def run_command():
+def run_command(wait_for_dag_run):
"""Fixture that provides a helper to run airflowctl commands."""
def _run_command(command: str, env_vars: dict, skip_login: bool = False)
-> str:
@@ -134,7 +193,7 @@ def run_command():
# Dag run to be terminal before running.
xcom_match = _XCOM_TARGET_PATTERN.match(command)
if xcom_match:
- _wait_for_dag_run_terminal_state(xcom_match.group(1),
xcom_match.group(2), env_vars, skip_login)
+ wait_for_dag_run(xcom_match.group(1), xcom_match.group(2))
console.print(f"[yellow]Running command: {command}")
@@ -189,10 +248,6 @@ def run_command():
return _run_command
-class _CtlTestState:
- docker_client: DockerClient | None = None
-
-
# Pytest hook to run at the start of the session
def pytest_sessionstart(session):
"""Install airflowctl at the very start of the pytest session."""
diff --git a/airflow-ctl-tests/tests/airflowctl_tests/constants.py
b/airflow-ctl-tests/tests/airflowctl_tests/constants.py
index 4ddf2636c7f..8a8bf812ac0 100644
--- a/airflow-ctl-tests/tests/airflowctl_tests/constants.py
+++ b/airflow-ctl-tests/tests/airflowctl_tests/constants.py
@@ -31,5 +31,14 @@ DOCKER_COMPOSE_FILE_PATH = (
AIRFLOW_ROOT_PATH / "airflow-core" / "docs" / "howto" / "docker-compose" /
"docker-compose.yaml"
)
-LOGIN_COMMAND = "auth login --username airflow --password airflow"
+API_USERNAME = "airflow"
+API_PASSWORD = "airflow"
+
+LOGIN_COMMAND = f"auth login --username {API_USERNAME} --password
{API_PASSWORD}"
LOGIN_OUTPUT = "Login successful! Welcome to airflowctl!"
+
+DAG_RUN_WAIT_TIMEOUT_ENV = "AIRFLOW_CTL_TEST_DAG_RUN_TIMEOUT"
+# Budget for a freshly booted compose stack to schedule and finish a Dag run:
the
+# scheduler still has to warm up, parse the examples and queue the tasks, none
of which
+# is a failure. Slower machines can raise it through the environment variable.
+DAG_RUN_WAIT_TIMEOUT = float(os.environ.get(DAG_RUN_WAIT_TIMEOUT_ENV, "300"))
diff --git a/airflow-ctl-tests/tests/airflowctl_tests/test_dag_run_wait.py
b/airflow-ctl-tests/tests/airflowctl_tests/test_dag_run_wait.py
new file mode 100644
index 00000000000..fe98e8ebb85
--- /dev/null
+++ b/airflow-ctl-tests/tests/airflowctl_tests/test_dag_run_wait.py
@@ -0,0 +1,222 @@
+# 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
+
+from itertools import islice
+
+import pytest
+import requests
+
+from airflowctl_tests import conftest
+
+
+class _JsonResponse:
+ def __init__(self, payload):
+ self._payload = payload
+
+ def raise_for_status(self):
+ return None
+
+ def json(self):
+ return self._payload
+
+
+def http_error(status_code: int) -> requests.exceptions.HTTPError:
+ response = requests.Response()
+ response.status_code = status_code
+ return requests.exceptions.HTTPError(str(status_code), response=response)
+
+
+def test_poll_delays_grow_and_cap():
+ assert list(islice(conftest._compute_poll_delays(), 8)) == [0.5, 1.0, 2.0,
4.0, 8.0, 10.0, 10.0, 10.0]
+
+
[email protected]("terminal_state", ["success", "failed"])
+def test_wait_returns_as_soon_as_the_dag_run_is_terminal(monkeypatch,
terminal_state):
+ states = iter(["queued", "running", terminal_state])
+ slept: list[float] = []
+ monkeypatch.setattr(conftest, "_find_dag_run_state", lambda dag_id,
dag_run_id: next(states))
+ monkeypatch.setattr(conftest.time, "sleep", slept.append)
+
+ conftest._wait_for_dag_run_terminal_state("example_bash_operator",
"manual__1", timeout=300)
+
+ assert slept == [0.5, 1.0]
+
+
+def test_wait_spends_its_whole_budget_and_no_more(monkeypatch):
+ now = [1000.0]
+ slept: list[float] = []
+
+ def sleep(seconds):
+ slept.append(seconds)
+ now[0] += seconds
+
+ monkeypatch.setattr(conftest, "_find_dag_run_state", lambda dag_id,
dag_run_id: "running")
+ monkeypatch.setattr(conftest, "_describe_task_instances", lambda dag_id,
dag_run_id: "runme_0=queued")
+ monkeypatch.setattr(conftest.time, "monotonic", lambda: now[0])
+ monkeypatch.setattr(conftest.time, "sleep", sleep)
+
+ with pytest.raises(TimeoutError):
+ conftest._wait_for_dag_run_terminal_state("example_bash_operator",
"manual__1", timeout=5)
+
+ # The last delay is clamped so the wait stops at the deadline instead of
overshooting it.
+ assert slept == [0.5, 1.0, 2.0, 1.5]
+
+
+def test_wait_timeout_reports_dag_run_and_task_instance_states(monkeypatch):
+ monkeypatch.setattr(conftest, "_find_dag_run_state", lambda dag_id,
dag_run_id: "running")
+ monkeypatch.setattr(
+ conftest, "_describe_task_instances", lambda dag_id, dag_run_id:
"runme_0=queued, runme_1=None"
+ )
+
+ with pytest.raises(TimeoutError) as exc_info:
+ conftest._wait_for_dag_run_terminal_state("example_bash_operator",
"manual__1", timeout=0)
+
+ message = str(exc_info.value)
+ assert "example_bash_operator/manual__1" in message
+ assert "Dag run state: running" in message
+ assert "runme_0=queued, runme_1=None" in message
+ assert conftest.DAG_RUN_WAIT_TIMEOUT_ENV in message
+
+
[email protected](
+ ("error", "retries"),
+ [
+ (requests.exceptions.ConnectionError("connection refused"), True),
+ (http_error(503), True),
+ (http_error(404), False),
+ (http_error(401), False),
+ ],
+ ids=["connection-error", "server-error", "missing-dag-run",
"rejected-token"],
+)
+def test_find_dag_run_state_retries_only_transient_failures(monkeypatch,
error, retries):
+ def raise_error(path):
+ raise error
+
+ monkeypatch.setattr(conftest, "_request_api", raise_error)
+
+ if retries:
+ assert conftest._find_dag_run_state("example_bash_operator",
"manual__1") is None
+ else:
+ with pytest.raises(requests.exceptions.HTTPError):
+ conftest._find_dag_run_state("example_bash_operator", "manual__1")
+
+
+def test_access_token_is_obtained_once_per_session(monkeypatch):
+ logins: list[tuple[str, str, str]] = []
+ monkeypatch.setattr(conftest._CtlTestState, "access_token", None)
+ monkeypatch.setattr(
+ conftest,
+ "generate_access_token",
+ lambda username, password, host: logins.append((username, password,
host)) or "token",
+ )
+
+ assert conftest._get_access_token() == "token"
+ assert conftest._get_access_token() == "token"
+ assert logins == [(conftest.API_USERNAME, conftest.API_PASSWORD,
conftest.DOCKER_COMPOSE_HOST_PORT)]
+
+
+def
test_request_api_targets_the_compose_stack_with_the_session_token(monkeypatch):
+ calls: list[tuple[str, dict]] = []
+ monkeypatch.setattr(conftest, "_get_access_token", lambda: "token")
+ monkeypatch.setattr(
+ conftest.requests,
+ "get",
+ lambda url, **kwargs: calls.append((url, kwargs)) or
_JsonResponse({"state": "success"}),
+ )
+
+ assert
conftest._request_api("dags/example_bash_operator/dagRuns/manual__1") ==
{"state": "success"}
+
+ url, kwargs = calls[0]
+ assert url == (
+
f"http://{conftest.DOCKER_COMPOSE_HOST_PORT}/api/v2/dags/example_bash_operator/dagRuns/manual__1";
+ )
+ assert kwargs["headers"] == {"Authorization": "Bearer token"}
+ assert kwargs["timeout"] == conftest._API_REQUEST_TIMEOUT
+
+
[email protected](
+ ("task_instances", "expected"),
+ [
+ (
+ [{"task_id": "runme_0", "state": "queued"}, {"task_id":
"run_this_last", "state": None}],
+ "runme_0=queued, run_this_last=None",
+ ),
+ ([], "none"),
+ ],
+ ids=["renders-each-task", "no-task-instances"],
+)
+def test_describe_task_instances_renders_every_task_state(monkeypatch,
task_instances, expected):
+ monkeypatch.setattr(conftest, "_request_api", lambda path:
{"task_instances": task_instances})
+
+ assert conftest._describe_task_instances("example_bash_operator",
"manual__1") == expected
+
+
+def test_describe_task_instances_reports_an_unreachable_api(monkeypatch):
+ def raise_connection_error(path):
+ raise requests.exceptions.ConnectionError("connection refused")
+
+ monkeypatch.setattr(conftest, "_request_api", raise_connection_error)
+
+ assert conftest._describe_task_instances("example_bash_operator",
"manual__1") == (
+ "unavailable (connection refused)"
+ )
+
+
+def test_waiter_spends_the_budget_once_per_dag_run(monkeypatch):
+ waited: list[tuple[str, str]] = []
+ monkeypatch.setattr(
+ conftest,
+ "_wait_for_dag_run_terminal_state",
+ lambda dag_id, dag_run_id: waited.append((dag_id, dag_run_id)),
+ )
+ wait = conftest._build_dag_run_waiter()
+
+ wait("example_bash_operator", "manual__1")
+ wait("example_bash_operator", "manual__1")
+ wait("example_bash_operator", "manual__2")
+
+ assert waited == [("example_bash_operator", "manual__1"),
("example_bash_operator", "manual__2")]
+
+
+def _build_timed_out_waiter(monkeypatch):
+ def raise_timeout(dag_id, dag_run_id):
+ raise TimeoutError("did not reach a terminal state")
+
+ monkeypatch.setattr(conftest, "_wait_for_dag_run_terminal_state",
raise_timeout)
+ wait = conftest._build_dag_run_waiter()
+ with pytest.raises(TimeoutError):
+ wait("example_bash_operator", "manual__1")
+ return wait
+
+
+def test_waiter_skips_later_commands_instead_of_timing_out_again(monkeypatch):
+ wait = _build_timed_out_waiter(monkeypatch)
+ monkeypatch.setattr(conftest, "_find_dag_run_state", lambda dag_id,
dag_run_id: "running")
+
+ with pytest.raises(pytest.skip.Exception, match="did not reach a terminal
state"):
+ wait("example_bash_operator", "manual__1")
+
+
+def
test_waiter_lets_later_commands_through_when_the_dag_run_finished_late(monkeypatch):
+ wait = _build_timed_out_waiter(monkeypatch)
+ monkeypatch.setattr(conftest, "_find_dag_run_state", lambda dag_id,
dag_run_id: "success")
+
+ try:
+ wait("example_bash_operator", "manual__1")
+ except pytest.skip.Exception:
+ pytest.fail("a Dag run that finished late must not skip the remaining
xcom commands")
diff --git a/contributing-docs/testing/airflow_ctl_tests.rst
b/contributing-docs/testing/airflow_ctl_tests.rst
index 4a678636210..f192a0c6379 100644
--- a/contributing-docs/testing/airflow_ctl_tests.rst
+++ b/contributing-docs/testing/airflow_ctl_tests.rst
@@ -45,3 +45,15 @@ Then, you can run the tests using the following command:
The above command runs the integration tests for ``airflowctl`` in the
Breeze environment.
Ensure that the correct tag or branch is checked out before executing the
tests.
+
+Waiting for Dag Runs
+--------------------
+
+The ``xcom`` commands can only run once their Dag run has finished, so the
tests wait for it.
+A freshly booted compose stack has to start the scheduler, queue the run and
execute the tasks,
+which on a slow machine can take longer than the default budget of 300 seconds:
+
+.. code-block:: bash
+
+ # Allow ten minutes for the Dag run to finish
+ AIRFLOW_CTL_TEST_DAG_RUN_TIMEOUT=600 breeze testing
airflow-ctl-integration-tests
