|
| 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 | +End-to-end test of a Dag declared entirely in TypeScript. |
| 19 | +
|
| 20 | +Run with:: |
| 21 | +
|
| 22 | + E2E_TEST_MODE=ts_sdk RUN_TS_SDK_NATIVE_DAG_TESTS=true \\ |
| 23 | + uv run --project airflow-e2e-tests pytest \\ |
| 24 | + tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_native_dag.py -xvs |
| 25 | +
|
| 26 | +Unlike ``test_ts_sdk_dag.py``, no Python file declares this Dag: the Dag processor asks the |
| 27 | +``airflow-ts-pack`` bundle to parse itself, and the bundle answers with the serialized Dag that |
| 28 | +``ts-sdk/example/src/native.ts`` built. The graph exercises every construct of the native surface |
| 29 | +rather than a linear chain -- a task group, a named fan-in, an order-only edge, a conditional, a |
| 30 | +multi-way branch, and a ``TriggerDagRunOperator`` a Python worker runs. |
| 31 | +
|
| 32 | +Gated behind ``RUN_TS_SDK_NATIVE_DAG_TESTS`` because it needs a Dag processor that can dispatch a |
| 33 | +parse request to a language coordinator. Until that lands, the bundle answers a request nothing |
| 34 | +sends, and the Dag never appears. |
| 35 | +""" |
| 36 | + |
| 37 | +from __future__ import annotations |
| 38 | + |
| 39 | +import os |
| 40 | +from dataclasses import dataclass |
| 41 | +from datetime import datetime, timezone |
| 42 | + |
| 43 | +import pytest |
| 44 | + |
| 45 | +from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient |
| 46 | + |
| 47 | +_RUN_NATIVE = os.environ.get("RUN_TS_SDK_NATIVE_DAG_TESTS", "").lower() in ("true", "1") |
| 48 | + |
| 49 | +# Parsing the bundle launches node before the first task is even scheduled, so |
| 50 | +# allow the same headroom the mixed-language suite does. |
| 51 | +_TS_TASK_TIMEOUT = 600 |
| 52 | + |
| 53 | +_DAG_ID = "typescript_native_example" |
| 54 | + |
| 55 | +# Read by the extract tasks; see ts-sdk/example/src/native.ts. |
| 56 | +_NORTH_ROWS_VARIABLE = "typescript_native_north_rows" |
| 57 | +_SOUTH_ROWS_VARIABLE = "typescript_native_south_rows" |
| 58 | + |
| 59 | +pytestmark = pytest.mark.skipif( |
| 60 | + not _RUN_NATIVE, |
| 61 | + reason="Needs a Dag processor that dispatches parse requests to a language coordinator " |
| 62 | + "(RUN_TS_SDK_NATIVE_DAG_TESTS)", |
| 63 | +) |
| 64 | + |
| 65 | + |
| 66 | +@dataclass |
| 67 | +class _CompletedRun: |
| 68 | + client: AirflowClient |
| 69 | + run_id: str |
| 70 | + state: str |
| 71 | + ti_states: dict[str, str] |
| 72 | + |
| 73 | + def xcom(self, task_id: str, key: str = "return_value"): |
| 74 | + return self.client.get_xcom_value(dag_id=_DAG_ID, task_id=task_id, run_id=self.run_id, key=key).get( |
| 75 | + "value" |
| 76 | + ) |
| 77 | + |
| 78 | + |
| 79 | +@pytest.fixture(scope="module") |
| 80 | +def completed_run() -> _CompletedRun: |
| 81 | + """Trigger the native Dag once with rows in both regions; every test reads that run.""" |
| 82 | + client = AirflowClient() |
| 83 | + # Both regions non-empty, so the conditional takes its `then` branch. |
| 84 | + for key, value in ((_NORTH_ROWS_VARIABLE, "3"), (_SOUTH_ROWS_VARIABLE, "2")): |
| 85 | + client.set_variable(key, value) |
| 86 | + |
| 87 | + resp = client.trigger_dag(_DAG_ID, json={"logical_date": datetime.now(timezone.utc).isoformat()}) |
| 88 | + run_id = resp["dag_run_id"] |
| 89 | + state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id, timeout=_TS_TASK_TIMEOUT) |
| 90 | + ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id) |
| 91 | + return _CompletedRun( |
| 92 | + client=client, |
| 93 | + run_id=run_id, |
| 94 | + state=state, |
| 95 | + ti_states={ti["task_id"]: ti.get("state") for ti in ti_resp.get("task_instances", [])}, |
| 96 | + ) |
| 97 | + |
| 98 | + |
| 99 | +def test_the_dag_the_bundle_parsed_is_registered(): |
| 100 | + """The Dag exists without any Python file declaring it.""" |
| 101 | + client = AirflowClient() |
| 102 | + dag = client.get_dag(_DAG_ID) |
| 103 | + |
| 104 | + assert dag["dag_id"] == _DAG_ID |
| 105 | + assert dag.get("timetable_summary") == "@daily" |
| 106 | + assert set(dag.get("tags") or []) >= {"typescript", "native"} |
| 107 | + |
| 108 | + |
| 109 | +def test_the_graph_carries_every_construct(): |
| 110 | + """Group prefixes, the fan-in, the branches and the trigger all survive parsing.""" |
| 111 | + tasks = AirflowClient().get_tasks(_DAG_ID) |
| 112 | + by_id = {task["task_id"]: task for task in tasks.get("tasks", [])} |
| 113 | + |
| 114 | + # The task group prefixed its members. |
| 115 | + assert {"extract.north", "extract.south"} <= set(by_id) |
| 116 | + # The named fan-in put both extract tasks upstream of summarize. |
| 117 | + assert set(by_id["summarize"]["upstream_task_ids"]) >= {"extract.north", "extract.south"} |
| 118 | + # The conditional and the branch reach their candidates. |
| 119 | + assert set(by_id["has_rows"]["downstream_task_ids"]) >= {"load_rows", "report_empty"} |
| 120 | + assert set(by_id["pick_cadence"]["downstream_task_ids"]) >= { |
| 121 | + "publish_daily", |
| 122 | + "publish_weekly", |
| 123 | + } |
| 124 | + # The order-only edge, which carries no value. |
| 125 | + assert "trigger_downstream" in by_id["cleanup"]["downstream_task_ids"] |
| 126 | + # And the one task Python runs rather than the Node coordinator. |
| 127 | + assert by_id["trigger_downstream"]["operator_name"] == "TriggerDagRunOperator" |
| 128 | + |
| 129 | + |
| 130 | +def test_dag_run_succeeded(completed_run: _CompletedRun): |
| 131 | + assert completed_run.state == "success", ( |
| 132 | + f"expected the run to succeed; got {completed_run.state!r}. task states: {completed_run.ti_states}" |
| 133 | + ) |
| 134 | + |
| 135 | + |
| 136 | +def test_the_taken_branches_ran_and_the_others_skipped(completed_run: _CompletedRun): |
| 137 | + """A branch is a run-time skip, so the states are what prove it worked.""" |
| 138 | + states = completed_run.ti_states |
| 139 | + |
| 140 | + # Both regions had rows, so the conditional followed `then`. |
| 141 | + assert states.get("load_rows") == "success" |
| 142 | + assert states.get("report_empty") == "skipped" |
| 143 | + |
| 144 | + # Exactly one cadence ran; which one depends on the run's own date. |
| 145 | + cadences = {states.get("publish_daily"), states.get("publish_weekly")} |
| 146 | + assert cadences == {"success", "skipped"}, f"expected one of each, got {cadences}" |
| 147 | + |
| 148 | + |
| 149 | +def test_every_other_task_succeeded(completed_run: _CompletedRun): |
| 150 | + always_run = [ |
| 151 | + "extract.north", |
| 152 | + "extract.south", |
| 153 | + "summarize", |
| 154 | + "has_rows", |
| 155 | + "pick_cadence", |
| 156 | + "cleanup", |
| 157 | + "trigger_downstream", |
| 158 | + ] |
| 159 | + for task_id in always_run: |
| 160 | + assert completed_run.ti_states.get(task_id) == "success", ( |
| 161 | + f"{task_id!r} did not succeed. all task states: {completed_run.ti_states}" |
| 162 | + ) |
| 163 | + |
| 164 | + |
| 165 | +def test_xcoms_flow_between_typescript_tasks(completed_run: _CompletedRun): |
| 166 | + """The fan-in's arguments reached the handler, and its own push came back.""" |
| 167 | + assert completed_run.xcom("extract.north") == {"region": "north", "rows": 3} |
| 168 | + assert completed_run.xcom("extract.south") == {"region": "south", "rows": 2} |
| 169 | + assert completed_run.xcom("summarize") == {"total": 5, "regions": 2} |
| 170 | + # Pushed under its own key by the summarize handler, and read back by load_rows. |
| 171 | + assert completed_run.xcom("summarize", key="region_total") == 5 |
| 172 | + assert completed_run.xcom("load_rows") == {"loaded": 5} |
| 173 | + |
| 174 | + |
| 175 | +def test_the_condition_reports_what_it_decided(completed_run: _CompletedRun): |
| 176 | + """A condition returns its boolean, and a branch the task id it chose.""" |
| 177 | + assert completed_run.xcom("has_rows") is True |
| 178 | + assert completed_run.xcom("pick_cadence") in ("publish_daily", "publish_weekly") |
0 commit comments