Skip to content

Commit 50fa403

Browse files
jason810496jscheffl
jason810496
and
jscheffl
authored
Add e2e test suite for Airflow event-driven DAGs with Apache Kafka (#64833)
* Add e2e test suite for Airflow event-driven DAGs with Apache Kafka * Add 'event_driven' option to E2E test mode in testing_commands.py * Address review feedback for event-driven E2E tests - Use timedelta(seconds=1) for retry_delay instead of bare int - Return exactly expected_count runs from _wait_for_consumer_dag_runs - Replace fixed 30s sleep with polling for Kafka consumer group registration - Pin --partitions 1 --replication-factor 1 for deterministic topic creation - Add kafka.yml pattern to selective checks file group - Remove container_name and unused JMX port from kafka.yml - Regenerate breeze CLI doc images * Fix AwaitTrigger and airflow e2e test * Add retry for un_pause_dag method * Add additional requirements for Kafka integration in event-driven setup * Add pytest output * fixup: Refactor comments and documentation for consistency in DAG terminology and update regex patterns for Kafka provider paths * fixup: Add entry in ci-amd.yml * fixup: Remove debug -sv flag * Apply suggestions from code review Co-authored-by: Jens Scheffler <95105677+jscheffl@users.noreply.github.com> * fixup: mount latest provider changes instead of installing from pypi --------- Co-authored-by: Jens Scheffler <95105677+jscheffl@users.noreply.github.com>
1 parent 7e22395 commit 50fa403

16 files changed

Lines changed: 701 additions & 5 deletions

File tree

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@ on: # yamllint disable-line rule:truthy
4848
description: "Whether to run OpenSearch remote logging e2e tests (true/false)"
4949
required: true
5050
type: string
51+
run-event-driven-e2e-tests:
52+
description: "Whether to run event driven e2e tests (true/false)"
53+
required: true
54+
type: string
5155
constraints-branch:
5256
description: "Branch used to construct constraints URL from."
5357
required: true
@@ -274,6 +278,18 @@ jobs:
274278
use-uv: ${{ inputs.use-uv }}
275279
e2e_test_mode: "xcom_object_storage"
276280

281+
test-e2e-integration-tests-event-driven:
282+
name: "Event driven tests with PROD image"
283+
uses: ./.github/workflows/airflow-e2e-tests.yml
284+
with:
285+
workflow-name: "Event driven e2e test"
286+
runners: ${{ inputs.runners }}
287+
platform: ${{ inputs.platform }}
288+
default-python-version: "${{ inputs.default-python-version }}"
289+
use-uv: ${{ inputs.use-uv }}
290+
e2e_test_mode: "event_driven"
291+
if: inputs.canary-run == 'true' || inputs.run-event-driven-e2e-tests == 'true'
292+
277293
test-ui-e2e-chromium:
278294
name: "Chromium UI e2e tests with PROD image"
279295
uses: ./.github/workflows/ui-e2e-tests.yml

‎.github/workflows/airflow-e2e-tests.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ on: # yamllint disable-line rule:truthy
4949
type: string
5050
required: true
5151
e2e_test_mode:
52-
description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, or xcom_object_storage" # yamllint disable-line rule:line-length
52+
description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, or event_driven" # yamllint disable-line rule:line-length
5353
type: string
5454
default: "basic"
5555

@@ -80,7 +80,7 @@ on: # yamllint disable-line rule:truthy
8080
type: string
8181
default: ""
8282
e2e_test_mode:
83-
description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, or xcom_object_storage" # yamllint disable-line rule:line-length
83+
description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, or event_driven" # yamllint disable-line rule:line-length
8484
type: string
8585
default: "basic"
8686

‎.github/workflows/ci-amd.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ jobs:
132132
run-remote-logging-elasticsearch-e2e-tests: ${{ steps.selective-checks.outputs.run-remote-logging-elasticsearch-e2e-tests }}
133133
run-remote-logging-opensearch-e2e-tests: ${{ steps.selective-checks.outputs.run-remote-logging-opensearch-e2e-tests }}
134134
run-remote-logging-s3-e2e-tests: ${{ steps.selective-checks.outputs.run-remote-logging-s3-e2e-tests }}
135+
run-event-driven-e2e-tests: ${{ steps.selective-checks.outputs.run-event-driven-e2e-tests }}
135136
run-system-tests: ${{ steps.selective-checks.outputs.run-system-tests }}
136137
run-task-sdk-tests: ${{ steps.selective-checks.outputs.run-task-sdk-tests }}
137138
run-task-sdk-integration-tests: ${{ steps.selective-checks.outputs.run-task-sdk-integration-tests }}
@@ -869,6 +870,7 @@ jobs:
869870
run-remote-logging-elasticsearch-e2e-tests: ${{ needs.build-info.outputs.run-remote-logging-elasticsearch-e2e-tests }}
870871
run-remote-logging-opensearch-e2e-tests: ${{ needs.build-info.outputs.run-remote-logging-opensearch-e2e-tests }}
871872
run-remote-logging-s3-e2e-tests: ${{ needs.build-info.outputs.run-remote-logging-s3-e2e-tests }}
873+
run-event-driven-e2e-tests: ${{ needs.build-info.outputs.run-event-driven-e2e-tests }}
872874
use-uv: ${{ needs.build-info.outputs.use-uv }}
873875
run-ui-e2e-tests: ${{ needs.build-info.outputs.run-ui-e2e-tests }}
874876
run-airflow-ctl-integration-tests: ${{ needs.build-info.outputs.run-airflow-ctl-integration-tests }}

‎.github/workflows/ci-arm.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ jobs:
122122
run-remote-logging-elasticsearch-e2e-tests: ${{ steps.selective-checks.outputs.run-remote-logging-elasticsearch-e2e-tests }}
123123
run-remote-logging-opensearch-e2e-tests: ${{ steps.selective-checks.outputs.run-remote-logging-opensearch-e2e-tests }}
124124
run-remote-logging-s3-e2e-tests: ${{ steps.selective-checks.outputs.run-remote-logging-s3-e2e-tests }}
125+
run-event-driven-e2e-tests: ${{ steps.selective-checks.outputs.run-event-driven-e2e-tests }}
125126
run-system-tests: ${{ steps.selective-checks.outputs.run-system-tests }}
126127
run-task-sdk-tests: ${{ steps.selective-checks.outputs.run-task-sdk-tests }}
127128
run-task-sdk-integration-tests: ${{ steps.selective-checks.outputs.run-task-sdk-integration-tests }}
@@ -859,6 +860,7 @@ jobs:
859860
run-remote-logging-elasticsearch-e2e-tests: ${{ needs.build-info.outputs.run-remote-logging-elasticsearch-e2e-tests }}
860861
run-remote-logging-opensearch-e2e-tests: ${{ needs.build-info.outputs.run-remote-logging-opensearch-e2e-tests }}
861862
run-remote-logging-s3-e2e-tests: ${{ needs.build-info.outputs.run-remote-logging-s3-e2e-tests }}
863+
run-event-driven-e2e-tests: ${{ needs.build-info.outputs.run-event-driven-e2e-tests }}
862864
use-uv: ${{ needs.build-info.outputs.use-uv }}
863865
run-ui-e2e-tests: ${{ needs.build-info.outputs.run-ui-e2e-tests }}
864866
run-airflow-ctl-integration-tests: ${{ needs.build-info.outputs.run-airflow-ctl-integration-tests }}

‎airflow-e2e-tests/docker/kafka.yml‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
---
18+
services:
19+
broker:
20+
image: confluentinc/cp-kafka:7.3.0
21+
labels:
22+
breeze.description: "Integration required for Kafka hooks."
23+
hostname: broker
24+
ports:
25+
- "9092:9092"
26+
environment:
27+
KAFKA_BROKER_ID: 1
28+
KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: >
29+
CONTROLLER:PLAINTEXT,
30+
PLAINTEXT:PLAINTEXT,
31+
PLAINTEXT_HOST:PLAINTEXT
32+
KAFKA_ADVERTISED_LISTENERS: >
33+
PLAINTEXT://broker:29092,
34+
PLAINTEXT_HOST://localhost:9092
35+
KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1
36+
KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS: 0
37+
KAFKA_TRANSACTION_STATE_LOG_MIN_ISR: 1
38+
KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: 1
39+
KAFKA_JMX_PORT: 9101
40+
KAFKA_JMX_HOSTNAME: localhost
41+
KAFKA_PROCESS_ROLES: 'broker,controller'
42+
KAFKA_NODE_ID: 1
43+
KAFKA_CONTROLLER_QUORUM_VOTERS: '1@broker:29093'
44+
KAFKA_LISTENERS: >
45+
PLAINTEXT://broker:29092,
46+
CONTROLLER://broker:29093,
47+
PLAINTEXT_HOST://0.0.0.0:9092
48+
KAFKA_INTER_BROKER_LISTENER_NAME: 'PLAINTEXT'
49+
KAFKA_CONTROLLER_LISTENER_NAMES: 'CONTROLLER'
50+
KAFKA_LOG_DIRS: '/tmp/kraft-combined-logs'
51+
volumes:
52+
- ./kafka/update_run.sh:/tmp/update_run.sh
53+
command: >
54+
bash -c 'if [ ! -f /tmp/update_run.sh ];
55+
then echo "ERROR: update_run.sh not mounted?"
56+
&& exit 1 ; else /tmp/update_run.sh && /etc/confluent/docker/run ; fi'

‎airflow-e2e-tests/docker/kafka/update_run.sh‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
#!/bin/sh
2+
3+
# Licensed to the Apache Software Foundation (ASF) under one
4+
# or more contributor license agreements. See the NOTICE file
5+
# distributed with this work for additional information
6+
# regarding copyright ownership. The ASF licenses this file
7+
# to you under the Apache License, Version 2.0 (the
8+
# "License"); you may not use this file except in compliance
9+
# with the License. You may obtain a copy of the License at
10+
#
11+
# http://www.apache.org/licenses/LICENSE-2.0
12+
#
13+
# Unless required by applicable law or agreed to in writing,
14+
# software distributed under the License is distributed on an
15+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16+
# KIND, either express or implied. See the License for the
17+
# specific language governing permissions and limitations
18+
# under the License.
19+
20+
# Docker workaround: Remove check for KAFKA_ZOOKEEPER_CONNECT parameter
21+
sed -i '/KAFKA_ZOOKEEPER_CONNECT/d' /etc/confluent/docker/configure
22+
23+
# Docker workaround: Ignore cub zk-ready
24+
sed -i 's/cub zk-ready/echo ignore zk-ready/' /etc/confluent/docker/ensure
25+
26+
# KRaft required step: Format the storage directory with a new cluster ID
27+
echo "kafka-storage format --ignore-formatted --cluster-id=$(kafka-storage random-uuid) -c /etc/kafka/kafka.properties" >> /etc/confluent/docker/ensure

‎airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,16 +27,20 @@
2727
from testcontainers.compose import DockerCompose
2828

2929
from airflow_e2e_tests.constants import (
30+
AIRFLOW_SERVICES_FOR_PROVIDER_MOUNT,
3031
AWS_INIT_PATH,
3132
DOCKER_COMPOSE_HOST_PORT,
3233
DOCKER_COMPOSE_PATH,
3334
DOCKER_IMAGE,
3435
E2E_DAGS_FOLDER,
3536
E2E_TEST_MODE,
3637
ELASTICSEARCH_PATH,
38+
KAFKA_DIR_PATH,
3739
LOCALSTACK_PATH,
3840
LOGS_FOLDER,
3941
OPENSEARCH_PATH,
42+
PROVIDERS_MOUNT_CONTAINER_PATH,
43+
PROVIDERS_ROOT_PATH,
4044
TEST_REPORT_FILE,
4145
XCOM_BUCKET,
4246
)
@@ -121,6 +125,94 @@ def _setup_opensearch_integration(dot_env_file, tmp_dir):
121125
os.environ["ENV_FILE_PATH"] = str(dot_env_file)
122126

123127

128+
def _copy_kafka_files(tmp_dir):
129+
"""Copy the Kafka compose file and broker init script into the temp directory."""
130+
copyfile(KAFKA_DIR_PATH.parent / "kafka.yml", tmp_dir / "kafka.yml")
131+
132+
kafka_dir = tmp_dir / "kafka"
133+
kafka_dir.mkdir()
134+
copyfile(KAFKA_DIR_PATH / "update_run.sh", kafka_dir / "update_run.sh")
135+
current_permissions = os.stat(kafka_dir / "update_run.sh").st_mode
136+
os.chmod(kafka_dir / "update_run.sh", current_permissions | 0o111)
137+
138+
139+
def _write_providers_mount_override(tmp_dir: Path, providers: list[str]) -> list[str]:
140+
"""Write a docker-compose override that bind-mounts in-tree provider sources.
141+
142+
Each entry in ``providers`` is a provider id with dot-separated path segments (e.g.
143+
``"apache.kafka"``). The host source ``providers/<dotted/as/slashes>`` is mounted
144+
read-only into every airflow service at ``<PROVIDERS_MOUNT_CONTAINER_PATH>/<dashed>``.
145+
Returns the list of in-container paths suitable for ``_PIP_ADDITIONAL_REQUIREMENTS``
146+
so pip installs the in-tree (latest, possibly unreleased) provider instead of the
147+
PyPI release.
148+
"""
149+
in_container_paths: list[str] = []
150+
volume_entries: list[str] = []
151+
for provider_id in providers:
152+
host_path = PROVIDERS_ROOT_PATH / provider_id.replace(".", "/")
153+
if not host_path.is_dir():
154+
raise RuntimeError(f"Provider source directory not found: {host_path}")
155+
container_path = f"{PROVIDERS_MOUNT_CONTAINER_PATH}/{provider_id.replace('.', '-')}"
156+
in_container_paths.append(container_path)
157+
volume_entries.append(f" - {host_path}:{container_path}:ro")
158+
159+
volumes_block = "\n".join(volume_entries)
160+
services_block = "\n".join(
161+
f" {svc}:\n volumes:\n{volumes_block}" for svc in AIRFLOW_SERVICES_FOR_PROVIDER_MOUNT
162+
)
163+
(tmp_dir / "providers-mount.yml").write_text(f"---\nservices:\n{services_block}\n")
164+
return in_container_paths
165+
166+
167+
def _setup_event_driven_integration(dot_env_file, tmp_dir):
168+
_copy_kafka_files(tmp_dir)
169+
170+
# Install kafka and common-messaging providers from the in-tree sources so the
171+
# test exercises the latest code even before a PyPI release is cut.
172+
provider_paths = _write_providers_mount_override(tmp_dir, ["apache.kafka", "common.messaging"])
173+
174+
kafka_conn = json.dumps(
175+
{
176+
"conn_type": "kafka",
177+
"extra": {
178+
"bootstrap.servers": "broker:29092",
179+
"group.id": "kafka_default_group",
180+
"security.protocol": "PLAINTEXT",
181+
"enable.auto.commit": False,
182+
"auto.offset.reset": "latest",
183+
},
184+
}
185+
)
186+
187+
dot_env_file.write_text(
188+
f"AIRFLOW_UID={os.getuid()}\n"
189+
f"AIRFLOW_CONN_KAFKA_DEFAULT='{kafka_conn}'\n"
190+
f"_PIP_ADDITIONAL_REQUIREMENTS={' '.join(provider_paths)}\n"
191+
)
192+
os.environ["ENV_FILE_PATH"] = str(dot_env_file)
193+
194+
195+
def _create_kafka_topics(compose_instance):
196+
"""Create Kafka topics required by the event-driven Dag."""
197+
for topic in ("fizz_buzz", "dlq"):
198+
compose_instance.exec_in_container(
199+
command=[
200+
"kafka-topics",
201+
"--bootstrap-server",
202+
"broker:29092",
203+
"--create",
204+
"--topic",
205+
topic,
206+
"--partitions",
207+
"1",
208+
"--replication-factor",
209+
"1",
210+
"--if-not-exists",
211+
],
212+
service_name="broker",
213+
)
214+
215+
124216
def _setup_xcom_object_storage_integration(dot_env_file, tmp_dir):
125217
_copy_localstack_files(tmp_dir)
126218

@@ -180,6 +272,9 @@ def spin_up_airflow_environment(tmp_path_factory: pytest.TempPathFactory):
180272
elif E2E_TEST_MODE == "xcom_object_storage":
181273
compose_file_names.append("localstack.yml")
182274
_setup_xcom_object_storage_integration(dot_env_file, tmp_dir)
275+
elif E2E_TEST_MODE == "event_driven":
276+
compose_file_names.extend(["kafka.yml", "providers-mount.yml"])
277+
_setup_event_driven_integration(dot_env_file, tmp_dir)
183278

184279
#
185280
# Please Do not use this Fernet key in any deployments! Please generate your own key.
@@ -204,6 +299,10 @@ def spin_up_airflow_environment(tmp_path_factory: pytest.TempPathFactory):
204299
command=["airflow", "dags", "reserialize"], service_name="airflow-dag-processor"
205300
)
206301

302+
if E2E_TEST_MODE == "event_driven":
303+
console.print("[yellow]Creating Kafka topics...")
304+
_create_kafka_topics(_E2ETestState.compose_instance)
305+
207306
except Exception:
208307
console.print("[red]Failed to start docker compose")
209308
if _E2ETestState.compose_instance:
@@ -265,6 +364,12 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int | pytest.ExitC
265364
_E2ETestState.compose_instance.stop()
266365

267366

367+
@pytest.fixture(scope="session")
368+
def compose_instance():
369+
"""Provide access to the running Docker Compose instance."""
370+
return _E2ETestState.compose_instance
371+
372+
268373
def generate_test_report(results):
269374
"""Generate test report with json summary."""
270375
report = {

‎airflow-e2e-tests/tests/airflow_e2e_tests/constants.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,18 @@
4747

4848
# s3 bucket name for XComObjectStorageBackend tests. This bucket will be created in the `init-aws.sh` script that is run as part of the LocalStack container initialization.
4949
XCOM_BUCKET = "test-xcom-objectstorage-backend"
50+
51+
KAFKA_DIR_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "kafka"
52+
53+
# Local provider sources are mounted into the airflow containers under this directory so
54+
# ``_PIP_ADDITIONAL_REQUIREMENTS`` can install the in-tree (latest, possibly unreleased)
55+
# provider rather than the published one from PyPI.
56+
PROVIDERS_ROOT_PATH = AIRFLOW_ROOT_PATH / "providers"
57+
PROVIDERS_MOUNT_CONTAINER_PATH = "/opt/airflow-providers"
58+
AIRFLOW_SERVICES_FOR_PROVIDER_MOUNT = (
59+
"airflow-apiserver",
60+
"airflow-scheduler",
61+
"airflow-dag-processor",
62+
"airflow-worker",
63+
"airflow-triggerer",
64+
)

0 commit comments

Comments
 (0)