parkhojeong commented on code in PR #64833:
URL: https://github.com/apache/airflow/pull/64833#discussion_r3177900187
##########
airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py:
##########
@@ -101,7 +101,30 @@ def _make_request(
response.raise_for_status()
return response.json()
+ def get_dag(self, dag_id: str):
+ return self._make_request(method="GET", endpoint=f"dags/{dag_id}")
+
+ def wait_for_dag(self, dag_id: str, timeout: int = 120, check_interval:
int = 3):
+ """Poll until *dag_id* is registered (parsed & serialized) and
reachable via the API.
+
+ On remote CI runners the DAG processor may need extra time to parse &
serialize
+ new DAG files, so calling ``un_pause_dag`` immediately can return a
404. Retry
Review Comment:
```suggestion
new Dag files, so calling ``un_pause_dag`` immediately can return a
404. Retry
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py:
##########
@@ -101,7 +101,30 @@ def _make_request(
response.raise_for_status()
return response.json()
+ def get_dag(self, dag_id: str):
+ return self._make_request(method="GET", endpoint=f"dags/{dag_id}")
+
+ def wait_for_dag(self, dag_id: str, timeout: int = 120, check_interval:
int = 3):
+ """Poll until *dag_id* is registered (parsed & serialized) and
reachable via the API.
+
+ On remote CI runners the DAG processor may need extra time to parse &
serialize
+ new DAG files, so calling ``un_pause_dag`` immediately can return a
404. Retry
+ until the DAG exists or *timeout* seconds elapse.
+ """
+ start = time.monotonic()
+ last_error: Exception | None = None
+ while time.monotonic() - start < timeout:
+ try:
+ return self.get_dag(dag_id)
+ except requests.HTTPError as exc:
+ if exc.response is None or exc.response.status_code != 404:
+ raise
+ last_error = exc
+ time.sleep(check_interval)
+ raise TimeoutError(f"DAG {dag_id} was not registered within
{timeout}s. Last error: {last_error}")
Review Comment:
```suggestion
raise TimeoutError(f"Dag {dag_id} was not registered within
{timeout}s. Last error: {last_error}")
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
+"""
+
+from __future__ import annotations
+
+import time
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+PRODUCER_DAG_ID = "event_driven_producer"
+CONSUMER_DAG_ID = "event_driven_consumer"
+# 8 valid orders + 1 malformed message
+EXPECTED_CONSUMER_RUNS = 9
+EXPECTED_FIZZ_BUZZ_OFFSET = 9
+EXPECTED_DLQ_OFFSET = 1
+
+
+def _parse_topic_offset(raw_output: str, topic: str) -> int:
+ """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and
return the offset."""
+ for line in raw_output.strip().splitlines():
+ parts = line.strip().split(":")
+ if len(parts) == 3 and parts[0] == topic:
+ return int(parts[2])
+ raise ValueError(f"Could not find offset for topic '{topic}' in
output:\n{raw_output}")
+
+
+class TestEventDrivenDag:
+ """Test the Kafka-based event-driven producer/consumer DAG pair."""
+
+ airflow_client = AirflowClient()
+
+ def _wait_for_kafka_consumer_group(self, compose_instance, timeout: int =
60, check_interval: int = 3):
+ """Poll until any Kafka consumer group is registered, indicating the
trigger is active."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-consumer-groups",
+ "--bootstrap-server",
+ "broker:29092",
+ "--list",
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ # Any non-empty group listing means the trigger's consumer has
registered
+ groups = [line.strip() for line in output.strip().splitlines() if
line.strip()]
+ if groups:
+ return
+ time.sleep(check_interval)
+ raise TimeoutError(f"No Kafka consumer group registered within
{timeout}s")
+
+ def _wait_for_consumer_dag_runs(
+ self, expected_count: int, timeout: int = 600, check_interval: int = 10
+ ) -> list[dict]:
+ """Poll until *expected_count* consumer DAG runs reach a terminal
state."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ terminal_runs = [r for r in runs if r["state"] in ("success",
"failed")]
+ if len(terminal_runs) >= expected_count:
+ # Return only the most recent expected_count runs to avoid
+ # assertion failures if extra runs exist from retries or prior
state.
+ terminal_runs.sort(
+ key=lambda r: (
+ r.get("end_date") or "",
+ r.get("start_date") or "",
+ r.get("logical_date") or "",
+ r.get("dag_run_id") or "",
+ ),
+ reverse=True,
+ )
+ return terminal_runs[:expected_count]
+ time.sleep(check_interval)
+
+ # Timed out — gather diagnostics
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ states = {r["dag_run_id"]: r["state"] for r in runs}
+ raise TimeoutError(
+ f"Expected {expected_count} terminal consumer DAG runs within
{timeout}s, "
+ f"but only found {len([r for r in runs if r['state'] in
('success', 'failed')])}. "
+ f"Run states: {states}"
+ )
+
+ def _get_topic_offset(self, compose_instance, topic: str) -> int:
+ """Return the current end-offset of *topic* via ``kafka-get-offsets``
inside the broker."""
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-get-offsets",
+ "--bootstrap-server",
+ "broker:29092",
+ "--topic",
+ topic,
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ return _parse_topic_offset(output, topic)
+
+ def _wait_for_topic_offset(
+ self, compose_instance, topic: str, expected_offset: int, timeout: int
= 30, check_interval: int = 2
+ ) -> int:
+ """Poll until *topic* reaches *expected_offset*, then return the
offset."""
+ start = time.monotonic()
+ offset = self._get_topic_offset(compose_instance, topic)
+ while offset < expected_offset and time.monotonic() - start < timeout:
+ time.sleep(check_interval)
+ offset = self._get_topic_offset(compose_instance, topic)
+ return offset
+
+ def _get_task_states(self, run_id: str) -> dict[str, str]:
+ """Return a mapping of task_id -> state for a consumer DAG run."""
+ response = self.airflow_client.get_task_instances(CONSUMER_DAG_ID,
run_id)
Review Comment:
```suggestion
"""Return a mapping of task_id -> state for a consumer Dag run."""
response = self.airflow_client.get_task_instances(CONSUMER_DAG_ID,
run_id)
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
+"""
+
+from __future__ import annotations
+
+import time
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+PRODUCER_DAG_ID = "event_driven_producer"
+CONSUMER_DAG_ID = "event_driven_consumer"
+# 8 valid orders + 1 malformed message
+EXPECTED_CONSUMER_RUNS = 9
+EXPECTED_FIZZ_BUZZ_OFFSET = 9
+EXPECTED_DLQ_OFFSET = 1
+
+
+def _parse_topic_offset(raw_output: str, topic: str) -> int:
+ """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and
return the offset."""
+ for line in raw_output.strip().splitlines():
+ parts = line.strip().split(":")
+ if len(parts) == 3 and parts[0] == topic:
+ return int(parts[2])
+ raise ValueError(f"Could not find offset for topic '{topic}' in
output:\n{raw_output}")
+
+
+class TestEventDrivenDag:
+ """Test the Kafka-based event-driven producer/consumer DAG pair."""
+
+ airflow_client = AirflowClient()
+
+ def _wait_for_kafka_consumer_group(self, compose_instance, timeout: int =
60, check_interval: int = 3):
+ """Poll until any Kafka consumer group is registered, indicating the
trigger is active."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-consumer-groups",
+ "--bootstrap-server",
+ "broker:29092",
+ "--list",
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ # Any non-empty group listing means the trigger's consumer has
registered
+ groups = [line.strip() for line in output.strip().splitlines() if
line.strip()]
+ if groups:
+ return
+ time.sleep(check_interval)
+ raise TimeoutError(f"No Kafka consumer group registered within
{timeout}s")
+
+ def _wait_for_consumer_dag_runs(
+ self, expected_count: int, timeout: int = 600, check_interval: int = 10
+ ) -> list[dict]:
+ """Poll until *expected_count* consumer DAG runs reach a terminal
state."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ terminal_runs = [r for r in runs if r["state"] in ("success",
"failed")]
+ if len(terminal_runs) >= expected_count:
+ # Return only the most recent expected_count runs to avoid
+ # assertion failures if extra runs exist from retries or prior
state.
+ terminal_runs.sort(
+ key=lambda r: (
+ r.get("end_date") or "",
+ r.get("start_date") or "",
+ r.get("logical_date") or "",
+ r.get("dag_run_id") or "",
+ ),
+ reverse=True,
+ )
+ return terminal_runs[:expected_count]
+ time.sleep(check_interval)
+
+ # Timed out — gather diagnostics
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ states = {r["dag_run_id"]: r["state"] for r in runs}
+ raise TimeoutError(
+ f"Expected {expected_count} terminal consumer DAG runs within
{timeout}s, "
+ f"but only found {len([r for r in runs if r['state'] in
('success', 'failed')])}. "
+ f"Run states: {states}"
+ )
+
+ def _get_topic_offset(self, compose_instance, topic: str) -> int:
+ """Return the current end-offset of *topic* via ``kafka-get-offsets``
inside the broker."""
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-get-offsets",
+ "--bootstrap-server",
+ "broker:29092",
+ "--topic",
+ topic,
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ return _parse_topic_offset(output, topic)
+
+ def _wait_for_topic_offset(
+ self, compose_instance, topic: str, expected_offset: int, timeout: int
= 30, check_interval: int = 2
+ ) -> int:
+ """Poll until *topic* reaches *expected_offset*, then return the
offset."""
+ start = time.monotonic()
+ offset = self._get_topic_offset(compose_instance, topic)
+ while offset < expected_offset and time.monotonic() - start < timeout:
+ time.sleep(check_interval)
+ offset = self._get_topic_offset(compose_instance, topic)
+ return offset
+
+ def _get_task_states(self, run_id: str) -> dict[str, str]:
+ """Return a mapping of task_id -> state for a consumer DAG run."""
+ response = self.airflow_client.get_task_instances(CONSUMER_DAG_ID,
run_id)
+ return {ti["task_id"]: ti["state"] for ti in
response["task_instances"]}
+
+ def test_producer_triggers_consumer_and_kafka_offsets(self,
compose_instance):
+ """Trigger the producer once and verify 9 consumer runs and Kafka
offsets.
+
+ Steps:
+ 1. Unpause the consumer DAG so the triggerer starts the AssetWatcher.
+ 2. Wait for the Kafka MessageQueueTrigger to begin polling.
+ 3. Trigger the producer DAG and wait for it to succeed.
+ 4. Wait for 9 consumer DAG runs to reach a terminal state.
+ 5. All 9 DAG runs succeed. Verify task-level behavior:
+ - 1 run has a failed ``process_message`` task and executes
``handle_dlq``.
+ - 8 runs succeed on ``process_message`` and skip ``handle_dlq``.
+ 6. Verify that the ``fizz_buzz`` topic has offset 9 (all messages
produced).
+ 7. Verify that the ``dlq`` topic has offset 1 (the malformed message).
+ """
+ # 1. Unpause consumer so the triggerer registers the AssetWatcher
+ self.airflow_client.un_pause_dag(CONSUMER_DAG_ID)
+
+ # 2. Wait for the triggerer to start the MessageQueueTrigger and
subscribe.
+ # The trigger uses poll_interval=1 and auto.offset.reset=latest so
it
+ # must be actively polling before the producer writes.
+ self._wait_for_kafka_consumer_group(compose_instance)
+
+ # 3. Trigger producer and wait for it to complete
+ producer_state =
self.airflow_client.trigger_dag_and_wait(PRODUCER_DAG_ID)
+ assert producer_state == "success", f"Producer DAG did not succeed.
Final state: {producer_state}"
+
+ # 4. Wait for all 9 consumer DAG runs
Review Comment:
```suggestion
# 4. Wait for all 9 consumer Dag runs
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
+"""
+
+from __future__ import annotations
+
+import time
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+PRODUCER_DAG_ID = "event_driven_producer"
+CONSUMER_DAG_ID = "event_driven_consumer"
+# 8 valid orders + 1 malformed message
+EXPECTED_CONSUMER_RUNS = 9
+EXPECTED_FIZZ_BUZZ_OFFSET = 9
+EXPECTED_DLQ_OFFSET = 1
+
+
+def _parse_topic_offset(raw_output: str, topic: str) -> int:
+ """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and
return the offset."""
+ for line in raw_output.strip().splitlines():
+ parts = line.strip().split(":")
+ if len(parts) == 3 and parts[0] == topic:
+ return int(parts[2])
+ raise ValueError(f"Could not find offset for topic '{topic}' in
output:\n{raw_output}")
+
+
+class TestEventDrivenDag:
+ """Test the Kafka-based event-driven producer/consumer DAG pair."""
+
+ airflow_client = AirflowClient()
+
+ def _wait_for_kafka_consumer_group(self, compose_instance, timeout: int =
60, check_interval: int = 3):
+ """Poll until any Kafka consumer group is registered, indicating the
trigger is active."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-consumer-groups",
+ "--bootstrap-server",
+ "broker:29092",
+ "--list",
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ # Any non-empty group listing means the trigger's consumer has
registered
+ groups = [line.strip() for line in output.strip().splitlines() if
line.strip()]
+ if groups:
+ return
+ time.sleep(check_interval)
+ raise TimeoutError(f"No Kafka consumer group registered within
{timeout}s")
+
+ def _wait_for_consumer_dag_runs(
+ self, expected_count: int, timeout: int = 600, check_interval: int = 10
+ ) -> list[dict]:
+ """Poll until *expected_count* consumer DAG runs reach a terminal
state."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ terminal_runs = [r for r in runs if r["state"] in ("success",
"failed")]
+ if len(terminal_runs) >= expected_count:
+ # Return only the most recent expected_count runs to avoid
+ # assertion failures if extra runs exist from retries or prior
state.
+ terminal_runs.sort(
+ key=lambda r: (
+ r.get("end_date") or "",
+ r.get("start_date") or "",
+ r.get("logical_date") or "",
+ r.get("dag_run_id") or "",
+ ),
+ reverse=True,
+ )
+ return terminal_runs[:expected_count]
+ time.sleep(check_interval)
+
+ # Timed out — gather diagnostics
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ states = {r["dag_run_id"]: r["state"] for r in runs}
+ raise TimeoutError(
+ f"Expected {expected_count} terminal consumer DAG runs within
{timeout}s, "
+ f"but only found {len([r for r in runs if r['state'] in
('success', 'failed')])}. "
+ f"Run states: {states}"
+ )
+
+ def _get_topic_offset(self, compose_instance, topic: str) -> int:
+ """Return the current end-offset of *topic* via ``kafka-get-offsets``
inside the broker."""
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-get-offsets",
+ "--bootstrap-server",
+ "broker:29092",
+ "--topic",
+ topic,
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ return _parse_topic_offset(output, topic)
+
+ def _wait_for_topic_offset(
+ self, compose_instance, topic: str, expected_offset: int, timeout: int
= 30, check_interval: int = 2
+ ) -> int:
+ """Poll until *topic* reaches *expected_offset*, then return the
offset."""
+ start = time.monotonic()
+ offset = self._get_topic_offset(compose_instance, topic)
+ while offset < expected_offset and time.monotonic() - start < timeout:
+ time.sleep(check_interval)
+ offset = self._get_topic_offset(compose_instance, topic)
+ return offset
+
+ def _get_task_states(self, run_id: str) -> dict[str, str]:
+ """Return a mapping of task_id -> state for a consumer DAG run."""
+ response = self.airflow_client.get_task_instances(CONSUMER_DAG_ID,
run_id)
+ return {ti["task_id"]: ti["state"] for ti in
response["task_instances"]}
+
+ def test_producer_triggers_consumer_and_kafka_offsets(self,
compose_instance):
+ """Trigger the producer once and verify 9 consumer runs and Kafka
offsets.
+
+ Steps:
+ 1. Unpause the consumer DAG so the triggerer starts the AssetWatcher.
+ 2. Wait for the Kafka MessageQueueTrigger to begin polling.
+ 3. Trigger the producer DAG and wait for it to succeed.
+ 4. Wait for 9 consumer DAG runs to reach a terminal state.
+ 5. All 9 DAG runs succeed. Verify task-level behavior:
+ - 1 run has a failed ``process_message`` task and executes
``handle_dlq``.
+ - 8 runs succeed on ``process_message`` and skip ``handle_dlq``.
+ 6. Verify that the ``fizz_buzz`` topic has offset 9 (all messages
produced).
+ 7. Verify that the ``dlq`` topic has offset 1 (the malformed message).
+ """
+ # 1. Unpause consumer so the triggerer registers the AssetWatcher
+ self.airflow_client.un_pause_dag(CONSUMER_DAG_ID)
+
+ # 2. Wait for the triggerer to start the MessageQueueTrigger and
subscribe.
+ # The trigger uses poll_interval=1 and auto.offset.reset=latest so
it
+ # must be actively polling before the producer writes.
+ self._wait_for_kafka_consumer_group(compose_instance)
+
+ # 3. Trigger producer and wait for it to complete
+ producer_state =
self.airflow_client.trigger_dag_and_wait(PRODUCER_DAG_ID)
+ assert producer_state == "success", f"Producer DAG did not succeed.
Final state: {producer_state}"
Review Comment:
```suggestion
assert producer_state == "success", f"Producer Dag did not succeed.
Final state: {producer_state}"
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py:
##########
@@ -101,7 +101,30 @@ def _make_request(
response.raise_for_status()
return response.json()
+ def get_dag(self, dag_id: str):
+ return self._make_request(method="GET", endpoint=f"dags/{dag_id}")
+
+ def wait_for_dag(self, dag_id: str, timeout: int = 120, check_interval:
int = 3):
+ """Poll until *dag_id* is registered (parsed & serialized) and
reachable via the API.
+
+ On remote CI runners the DAG processor may need extra time to parse &
serialize
+ new DAG files, so calling ``un_pause_dag`` immediately can return a
404. Retry
+ until the DAG exists or *timeout* seconds elapse.
Review Comment:
```suggestion
until the Dag exists or *timeout* seconds elapse.
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py:
##########
@@ -154,6 +177,13 @@ def get_task_instances(self, dag_id: str, run_id: str):
endpoint=f"dags/{dag_id}/dagRuns/{run_id}/taskInstances",
)
+ def list_dag_runs(self, dag_id: str, limit: int = 100):
+ """List DAG runs for a given DAG."""
Review Comment:
```suggestion
"""List Dag runs for a given DAG."""
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py:
##########
@@ -121,6 +122,63 @@ def _setup_opensearch_integration(dot_env_file, tmp_dir):
os.environ["ENV_FILE_PATH"] = str(dot_env_file)
+def _copy_kafka_files(tmp_dir):
+ """Copy Kafka compose file, init script, and provider source into the temp
directory."""
+ copyfile(KAFKA_DIR_PATH.parent / "kafka.yml", tmp_dir / "kafka.yml")
+
+ kafka_dir = tmp_dir / "kafka"
+ kafka_dir.mkdir()
+ copyfile(KAFKA_DIR_PATH / "update_run.sh", kafka_dir / "update_run.sh")
+ current_permissions = os.stat(kafka_dir / "update_run.sh").st_mode
+ os.chmod(kafka_dir / "update_run.sh", current_permissions | 0o111)
+
+
+def _setup_event_driven_integration(dot_env_file, tmp_dir):
+ _copy_kafka_files(tmp_dir)
+
+ kafka_conn = json.dumps(
+ {
+ "conn_type": "general",
+ "extra": {
+ "bootstrap.servers": "broker:29092",
+ "group.id": "kafka_default_group",
+ "security.protocol": "PLAINTEXT",
+ "enable.auto.commit": False,
+ "auto.offset.reset": "latest",
+ },
+ }
+ )
+
+ dot_env_file.write_text(
+ f"AIRFLOW_UID={os.getuid()}\n"
+ f"AIRFLOW_CONN_KAFKA_DEFAULT='{kafka_conn}'\n"
+ "_PIP_ADDITIONAL_REQUIREMENTS="
+ "apache-airflow-providers-apache-kafka
apache-airflow-providers-common-messaging\n"
+ )
+ os.environ["ENV_FILE_PATH"] = str(dot_env_file)
+
+
+def _create_kafka_topics(compose_instance):
+ """Create Kafka topics required by the event-driven DAG."""
Review Comment:
```suggestion
"""Create Kafka topics required by the event-driven Dag."""
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
+"""
+
+from __future__ import annotations
+
+import time
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+PRODUCER_DAG_ID = "event_driven_producer"
+CONSUMER_DAG_ID = "event_driven_consumer"
+# 8 valid orders + 1 malformed message
+EXPECTED_CONSUMER_RUNS = 9
+EXPECTED_FIZZ_BUZZ_OFFSET = 9
+EXPECTED_DLQ_OFFSET = 1
+
+
+def _parse_topic_offset(raw_output: str, topic: str) -> int:
+ """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and
return the offset."""
+ for line in raw_output.strip().splitlines():
+ parts = line.strip().split(":")
+ if len(parts) == 3 and parts[0] == topic:
+ return int(parts[2])
+ raise ValueError(f"Could not find offset for topic '{topic}' in
output:\n{raw_output}")
+
+
+class TestEventDrivenDag:
+ """Test the Kafka-based event-driven producer/consumer DAG pair."""
Review Comment:
```suggestion
"""Test the Kafka-based event-driven producer/consumer Dag pair."""
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
Review Comment:
```suggestion
"""E2E tests for the event-driven Dag pattern using Apache Kafka.
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
Review Comment:
```suggestion
The producer Dag sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
Kafka topic. The consumer Dag is scheduled on an Asset with a Kafka
AssetWatcher,
so each message triggers a separate consumer Dag run (9 total).
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
+"""
+
+from __future__ import annotations
+
+import time
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+PRODUCER_DAG_ID = "event_driven_producer"
+CONSUMER_DAG_ID = "event_driven_consumer"
+# 8 valid orders + 1 malformed message
+EXPECTED_CONSUMER_RUNS = 9
+EXPECTED_FIZZ_BUZZ_OFFSET = 9
+EXPECTED_DLQ_OFFSET = 1
+
+
+def _parse_topic_offset(raw_output: str, topic: str) -> int:
+ """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and
return the offset."""
+ for line in raw_output.strip().splitlines():
+ parts = line.strip().split(":")
+ if len(parts) == 3 and parts[0] == topic:
+ return int(parts[2])
+ raise ValueError(f"Could not find offset for topic '{topic}' in
output:\n{raw_output}")
+
+
+class TestEventDrivenDag:
+ """Test the Kafka-based event-driven producer/consumer DAG pair."""
+
+ airflow_client = AirflowClient()
+
+ def _wait_for_kafka_consumer_group(self, compose_instance, timeout: int =
60, check_interval: int = 3):
+ """Poll until any Kafka consumer group is registered, indicating the
trigger is active."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-consumer-groups",
+ "--bootstrap-server",
+ "broker:29092",
+ "--list",
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ # Any non-empty group listing means the trigger's consumer has
registered
+ groups = [line.strip() for line in output.strip().splitlines() if
line.strip()]
+ if groups:
+ return
+ time.sleep(check_interval)
+ raise TimeoutError(f"No Kafka consumer group registered within
{timeout}s")
+
+ def _wait_for_consumer_dag_runs(
+ self, expected_count: int, timeout: int = 600, check_interval: int = 10
+ ) -> list[dict]:
+ """Poll until *expected_count* consumer DAG runs reach a terminal
state."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ terminal_runs = [r for r in runs if r["state"] in ("success",
"failed")]
+ if len(terminal_runs) >= expected_count:
+ # Return only the most recent expected_count runs to avoid
+ # assertion failures if extra runs exist from retries or prior
state.
+ terminal_runs.sort(
+ key=lambda r: (
+ r.get("end_date") or "",
+ r.get("start_date") or "",
+ r.get("logical_date") or "",
+ r.get("dag_run_id") or "",
+ ),
+ reverse=True,
+ )
+ return terminal_runs[:expected_count]
+ time.sleep(check_interval)
+
+ # Timed out — gather diagnostics
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ states = {r["dag_run_id"]: r["state"] for r in runs}
+ raise TimeoutError(
+ f"Expected {expected_count} terminal consumer DAG runs within
{timeout}s, "
+ f"but only found {len([r for r in runs if r['state'] in
('success', 'failed')])}. "
+ f"Run states: {states}"
+ )
+
+ def _get_topic_offset(self, compose_instance, topic: str) -> int:
+ """Return the current end-offset of *topic* via ``kafka-get-offsets``
inside the broker."""
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-get-offsets",
+ "--bootstrap-server",
+ "broker:29092",
+ "--topic",
+ topic,
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ return _parse_topic_offset(output, topic)
+
+ def _wait_for_topic_offset(
+ self, compose_instance, topic: str, expected_offset: int, timeout: int
= 30, check_interval: int = 2
+ ) -> int:
+ """Poll until *topic* reaches *expected_offset*, then return the
offset."""
+ start = time.monotonic()
+ offset = self._get_topic_offset(compose_instance, topic)
+ while offset < expected_offset and time.monotonic() - start < timeout:
+ time.sleep(check_interval)
+ offset = self._get_topic_offset(compose_instance, topic)
+ return offset
+
+ def _get_task_states(self, run_id: str) -> dict[str, str]:
+ """Return a mapping of task_id -> state for a consumer DAG run."""
+ response = self.airflow_client.get_task_instances(CONSUMER_DAG_ID,
run_id)
+ return {ti["task_id"]: ti["state"] for ti in
response["task_instances"]}
+
+ def test_producer_triggers_consumer_and_kafka_offsets(self,
compose_instance):
+ """Trigger the producer once and verify 9 consumer runs and Kafka
offsets.
+
+ Steps:
+ 1. Unpause the consumer DAG so the triggerer starts the AssetWatcher.
+ 2. Wait for the Kafka MessageQueueTrigger to begin polling.
+ 3. Trigger the producer DAG and wait for it to succeed.
+ 4. Wait for 9 consumer DAG runs to reach a terminal state.
+ 5. All 9 DAG runs succeed. Verify task-level behavior:
Review Comment:
```suggestion
1. Unpause the consumer Dag so the triggerer starts the AssetWatcher.
2. Wait for the Kafka MessageQueueTrigger to begin polling.
3. Trigger the producer Dag and wait for it to succeed.
4. Wait for 9 consumer Dag runs to reach a terminal state.
5. All 9 Dag runs succeed. Verify task-level behavior:
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
+"""
+
+from __future__ import annotations
+
+import time
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+PRODUCER_DAG_ID = "event_driven_producer"
+CONSUMER_DAG_ID = "event_driven_consumer"
+# 8 valid orders + 1 malformed message
+EXPECTED_CONSUMER_RUNS = 9
+EXPECTED_FIZZ_BUZZ_OFFSET = 9
+EXPECTED_DLQ_OFFSET = 1
+
+
+def _parse_topic_offset(raw_output: str, topic: str) -> int:
+ """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and
return the offset."""
+ for line in raw_output.strip().splitlines():
+ parts = line.strip().split(":")
+ if len(parts) == 3 and parts[0] == topic:
+ return int(parts[2])
+ raise ValueError(f"Could not find offset for topic '{topic}' in
output:\n{raw_output}")
+
+
+class TestEventDrivenDag:
+ """Test the Kafka-based event-driven producer/consumer DAG pair."""
+
+ airflow_client = AirflowClient()
+
+ def _wait_for_kafka_consumer_group(self, compose_instance, timeout: int =
60, check_interval: int = 3):
+ """Poll until any Kafka consumer group is registered, indicating the
trigger is active."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-consumer-groups",
+ "--bootstrap-server",
+ "broker:29092",
+ "--list",
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ # Any non-empty group listing means the trigger's consumer has
registered
+ groups = [line.strip() for line in output.strip().splitlines() if
line.strip()]
+ if groups:
+ return
+ time.sleep(check_interval)
+ raise TimeoutError(f"No Kafka consumer group registered within
{timeout}s")
+
+ def _wait_for_consumer_dag_runs(
+ self, expected_count: int, timeout: int = 600, check_interval: int = 10
+ ) -> list[dict]:
+ """Poll until *expected_count* consumer DAG runs reach a terminal
state."""
Review Comment:
```suggestion
"""Poll until *expected_count* consumer Dag runs reach a terminal
state."""
```
##########
airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py:
##########
@@ -0,0 +1,229 @@
+# 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.
+"""E2E tests for the event-driven DAG pattern using Apache Kafka.
+
+The producer DAG sends 8 valid orders + 1 malformed message to the
``fizz_buzz``
+Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka
AssetWatcher,
+so each message triggers a separate consumer DAG run (9 total).
+"""
+
+from __future__ import annotations
+
+import time
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+PRODUCER_DAG_ID = "event_driven_producer"
+CONSUMER_DAG_ID = "event_driven_consumer"
+# 8 valid orders + 1 malformed message
+EXPECTED_CONSUMER_RUNS = 9
+EXPECTED_FIZZ_BUZZ_OFFSET = 9
+EXPECTED_DLQ_OFFSET = 1
+
+
+def _parse_topic_offset(raw_output: str, topic: str) -> int:
+ """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and
return the offset."""
+ for line in raw_output.strip().splitlines():
+ parts = line.strip().split(":")
+ if len(parts) == 3 and parts[0] == topic:
+ return int(parts[2])
+ raise ValueError(f"Could not find offset for topic '{topic}' in
output:\n{raw_output}")
+
+
+class TestEventDrivenDag:
+ """Test the Kafka-based event-driven producer/consumer DAG pair."""
+
+ airflow_client = AirflowClient()
+
+ def _wait_for_kafka_consumer_group(self, compose_instance, timeout: int =
60, check_interval: int = 3):
+ """Poll until any Kafka consumer group is registered, indicating the
trigger is active."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ stdout, _, _ = compose_instance.exec_in_container(
+ command=[
+ "kafka-consumer-groups",
+ "--bootstrap-server",
+ "broker:29092",
+ "--list",
+ ],
+ service_name="broker",
+ )
+ output = stdout.decode() if isinstance(stdout, bytes) else stdout
+ # Any non-empty group listing means the trigger's consumer has
registered
+ groups = [line.strip() for line in output.strip().splitlines() if
line.strip()]
+ if groups:
+ return
+ time.sleep(check_interval)
+ raise TimeoutError(f"No Kafka consumer group registered within
{timeout}s")
+
+ def _wait_for_consumer_dag_runs(
+ self, expected_count: int, timeout: int = 600, check_interval: int = 10
+ ) -> list[dict]:
+ """Poll until *expected_count* consumer DAG runs reach a terminal
state."""
+ start = time.monotonic()
+ while time.monotonic() - start < timeout:
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ terminal_runs = [r for r in runs if r["state"] in ("success",
"failed")]
+ if len(terminal_runs) >= expected_count:
+ # Return only the most recent expected_count runs to avoid
+ # assertion failures if extra runs exist from retries or prior
state.
+ terminal_runs.sort(
+ key=lambda r: (
+ r.get("end_date") or "",
+ r.get("start_date") or "",
+ r.get("logical_date") or "",
+ r.get("dag_run_id") or "",
+ ),
+ reverse=True,
+ )
+ return terminal_runs[:expected_count]
+ time.sleep(check_interval)
+
+ # Timed out — gather diagnostics
+ response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID)
+ runs = response.get("dag_runs", [])
+ states = {r["dag_run_id"]: r["state"] for r in runs}
+ raise TimeoutError(
+ f"Expected {expected_count} terminal consumer DAG runs within
{timeout}s, "
Review Comment:
```suggestion
f"Expected {expected_count} terminal consumer Dag runs within
{timeout}s, "
```
--
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]
