jason810496 commented on code in PR #63185:
URL: https://github.com/apache/airflow/pull/63185#discussion_r3540810296
##########
airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py:
##########
@@ -476,3 +489,1244 @@ def test_get_all_bundle_names():
# the naming suffix instead of pinning an exact list.
extra = [n for n in bundle_names if n not in {"dags-folder",
"example_dags"}]
assert all(n.endswith("-example-dags") for n in extra)
+
+
[email protected]
+def clear_dags_and_bundles():
+ clear_db_dags()
+ clear_db_dag_bundles()
+ yield
+ clear_db_dags()
+ clear_db_dag_bundles()
+
+
+def _add_dag(session, dag_id: str, bundle_name: str) -> DagModel:
+ dag = DagModel(dag_id=dag_id, bundle_name=bundle_name,
fileloc=f"/tmp/{dag_id}.py")
+ session.add(dag)
+ session.flush()
+ return dag
+
+
+class TestBestBundleForFileloc:
+ """Tests for ``_best_bundle_for_fileloc`` path matching."""
+
+ def test_returns_relative_path_for_match(self) -> None:
+ assert _best_bundle_for_fileloc("/dags/team_x/dag.py", {"team-x":
Path("/dags/team_x")}) == (
+ "team-x",
+ "dag.py",
+ )
+
+ def test_returns_none_for_no_match(self) -> None:
+ assert _best_bundle_for_fileloc("/elsewhere/dag.py", {"team-x":
Path("/dags/team_x")}) is None
+
+ def test_returns_none_for_empty_paths(self) -> None:
+ assert _best_bundle_for_fileloc("/dags/dag.py", {}) is None
+
+ def test_normalises_redundant_separators_and_dots(self) -> None:
+ # ``Path`` itself collapses ``//`` and ``.`` segments on construction,
+ # so the helper matches without any explicit normalization.
+ assert _best_bundle_for_fileloc("/dags//team_x/./dag.py", {"team-x":
Path("/dags/team_x")}) == (
+ "team-x",
+ "dag.py",
+ )
+
+
[email protected]_test
+class TestReassignDagsWithUnconfiguredBundles:
+ """Tests for DagBundlesManager.reassign_dags_with_unconfigured_bundles."""
+
+ def _manager_with_bundle_names(self, names: list[str]) ->
DagBundlesManager:
+ """Return a manager whose ``_resolve_active_bundle_paths`` reports
*names* with non-matching paths.
+
+ The fake paths are absolute roots that cannot contain any
+ ``/tmp/{dag_id}.py`` test fileloc, so rows are routed as unmatched
+ unless a test explicitly arranges otherwise.
+ """
+ with patch.dict(
+ os.environ,
+ {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST":
json.dumps(BASIC_BUNDLE_CONFIG)},
+ ):
+ manager = DagBundlesManager()
+ paths = {name: Path(f"/__unmatched_for_test__/{name}") for name in
names}
+ manager._resolve_active_bundle_paths = lambda *, session: paths #
type: ignore[method-assign]
+ return manager
+
+ def test_no_configured_bundles_is_noop(self, clear_dags_and_bundles,
session):
+ """Return 0 without raising when no bundles are configured."""
+ manager = self._manager_with_bundle_names([])
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+
+ def test_already_configured_is_noop(self, clear_dags_and_bundles, session)
-> None:
+ """No reassignment when every Dag already points at a configured
bundle."""
+ bundle = DagBundleModel(name="bundle-a")
+ bundle.active = True
+ session.add(bundle)
+ session.flush()
+ _add_dag(session, "dag-1", "bundle-a")
+ session.commit()
+
+ manager = self._manager_with_bundle_names(["bundle-a"])
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+ assert session.get(DagModel, "dag-1").bundle_name == "bundle-a"
+
+ def test_unmatched_fileloc_leaves_row_untouched(self,
clear_dags_and_bundles, session, caplog) -> None:
+ """Rows whose fileloc has no configured bundle path keep their
original bundle.
+
+ The fallback that wrote ``bundle_name`` without a verified
+ ``relative_fileloc`` produced active-but-un-runnable rows, so the
+ helper now leaves such rows on their unconfigured bundle and emits a
+ single warning naming the count.
+ """
+ active = DagBundleModel(name="active")
+ active.active = True
+ removed = DagBundleModel(name="removed-bundle")
+ removed.active = False
+ session.add(active)
+ session.add(removed)
+ session.flush()
+ _add_dag(session, "dag-1", "removed-bundle")
+ _add_dag(session, "dag-2", "removed-bundle")
+ session.commit()
+
+ manager = self._manager_with_bundle_names(["active"])
+ with caplog.at_level("WARNING",
logger="airflow.dag_processing.bundles.manager.DagBundlesManager"):
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+
+ session.expire_all()
+ for dag_id in ("dag-1", "dag-2"):
+ row = session.get(DagModel, dag_id)
+ assert row.bundle_name == "removed-bundle"
+ assert row.relative_fileloc is None
+ assert any("Skipped 2 legacy Dag(s)" in record.message for record in
caplog.records)
+
+ def test_row_with_populated_relative_fileloc_is_left_alone(self,
clear_dags_and_bundles, session) -> None:
+ """A 3.x row whose bundle is no longer configured must keep its bundle
assignment.
+
+ Only rows the 0082 migration touched (``relative_fileloc IS NULL``) are
+ candidates for reassignment. Rows that already carry a relative path
+ were written by 3.x serialization and must be left on their bundle so
+ the regular stale-Dag deactivation path can handle a removed bundle.
+ """
+ active_bundle = DagBundleModel(name="active")
+ active_bundle.active = True
+ removed_bundle = DagBundleModel(name="removed-bundle")
+ removed_bundle.active = False
+ session.add(active_bundle)
+ session.add(removed_bundle)
+ session.flush()
+
+ dag = DagModel(
+ dag_id="dag-1",
+ bundle_name="removed-bundle",
+ fileloc="/tmp/dag-1.py",
+ )
+ dag.relative_fileloc = "dag-1.py"
+ dag.bundle_version = "abc123"
+ session.add(dag)
+ session.flush()
+ session.commit()
+
+ manager = self._manager_with_bundle_names(["active"])
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+
+ session.expire_all()
+ refreshed = session.get(DagModel, "dag-1")
+ assert refreshed.bundle_name == "removed-bundle"
+ assert refreshed.relative_fileloc == "dag-1.py"
+ assert refreshed.bundle_version == "abc123"
+
+ def test_row_with_existing_dag_version_is_left_alone(self,
clear_dags_and_bundles, session) -> None:
+ """A Dag with any DagVersion row must not be reassigned.
+
+ Defended at two layers: (1) the global DagVersion-existence fast-skip
+ short-circuits before any work, and (2) the per-row predicate
+ ``NOT EXISTS DagVersion`` excludes the row even if the fast-skip is
+ bypassed. The parse path is the source of truth for any Dag with a
+ DagVersion -- touching only the DagModel would leave the DagVersion
+ stale, and scheduler/executor paths prefer DagVersion.bundle_name
+ when building task workloads.
+ """
+ from airflow.models.dag_version import DagVersion
Review Comment:
I parameterized the test and moved the import to the top level in
https://github.com/apache/airflow/pull/63185/changes/a1eedb2f0ce1001b268773db063f6996cbce7e03.
--
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]
