From 1dc4cfcff25ed438fd999428f6a54a247752d045 Mon Sep 17 00:00:00 2001 From: Jed Cunningham Date: Mon, 24 Aug 2026 07:38:24 -0600 Subject: [PATCH] Clean up comments and tests for the Dag version inflation checker The comments described how the code used to behave rather than what it does now, and the removed test covered the same case as the one above it. --- .../utils/dag_version_inflation_checker.py | 2 - .../test_dag_version_inflation_checker.py | 45 ++++--------------- 2 files changed, 8 insertions(+), 39 deletions(-) diff --git a/airflow-core/src/airflow/utils/dag_version_inflation_checker.py b/airflow-core/src/airflow/utils/dag_version_inflation_checker.py index 06fa96a7faae0..e9b415b255f1b 100644 --- a/airflow-core/src/airflow/utils/dag_version_inflation_checker.py +++ b/airflow-core/src/airflow/utils/dag_version_inflation_checker.py @@ -513,8 +513,6 @@ def visit_With(self, node: ast.With): for body in node.body: self.visit(body) - # Only exit Dag with block if we entered it; unconditional exit would prematurely - # reset the context when a non-DAG with-statement is nested inside a DAG with block. if is_with_dag_context: self.dag_detector.exit_dag_context() diff --git a/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py b/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py index 035375c01f348..91c468338ba6a 100644 --- a/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py +++ b/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py @@ -814,53 +814,25 @@ def test_task(): assert len(warnings) == 1 def test_nested_non_dag_with_does_not_exit_dag_context(self): - """Regression test: a non-DAG with-statement nested inside a DAG with-block must not - prematurely exit the DAG context. - - Before the fix, visit_With() called exit_dag_context() unconditionally regardless of - whether a DAG context was entered. This meant any nested non-DAG with-statement (e.g., - ``with open(...) as f``) would reset is_in_dag_context=False, causing tasks that appear - AFTER the nested with-statement (but still inside the DAG with-block) to be missed. - """ + """A non-Dag with-statement nested inside a Dag with-block must not exit the Dag context.""" code = """ from airflow import DAG from datetime import datetime from airflow.operators.bash import BashOperator with DAG('my_dag') as dag: - with open('some_file.txt') as f: # non-DAG with — must not exit DAG context + with open('some_file.txt') as f: data = f.read() - # This task is still inside the DAG block and uses datetime.now() — must be flagged t1 = BashOperator( task_id='test', - bash_command=str(datetime.now()), + bash_command=str(datetime.now()), # !problem ) -""" - warnings = self._check_code(code) - assert len(warnings) == 1, ( - "BashOperator after a nested non-DAG with-statement must still be detected " - "as inside the DAG context and flagged for using datetime.now()" - ) - - def test_multiple_nested_non_dag_withs_do_not_exit_dag_context(self): - """Multiple successive non-DAG with-statements inside a DAG block must not exit the context.""" - code = """ -from airflow import DAG -from datetime import datetime -from airflow.operators.bash import BashOperator - -with DAG('my_dag') as dag: - with open('a.txt') as f1: - pass - with open('b.txt') as f2: - pass - t1 = BashOperator(task_id='t', bash_command=str(datetime.now())) """ warnings = self._check_code(code) assert len(warnings) == 1 def test_task_inside_nested_non_dag_with_is_still_flagged(self): - """A task constructed inside a nested non-DAG with-block is still in DAG context.""" + """A task constructed inside a nested non-Dag with-block is still in Dag context.""" code = """ from airflow import DAG from datetime import datetime @@ -868,13 +840,13 @@ def test_task_inside_nested_non_dag_with_is_still_flagged(self): with DAG('my_dag') as dag: with open('a.txt') as f: - t1 = BashOperator(task_id='t', bash_command=str(datetime.now())) + t1 = BashOperator(task_id='t', bash_command=str(datetime.now())) # !problem """ warnings = self._check_code(code) assert len(warnings) == 1 - def test_task_outside_dag_with_not_flagged_after_nested_fix(self): - """Tasks genuinely outside any DAG block must not be flagged — the fix must not over-flag.""" + def test_task_outside_dag_with_not_flagged(self): + """Tasks genuinely outside any Dag block must not be flagged.""" code = """ from airflow import DAG from datetime import datetime @@ -884,8 +856,7 @@ def test_task_outside_dag_with_not_flagged_after_nested_fix(self): with open('a.txt') as f: pass -# This task is outside the DAG with block — should NOT be flagged t_outside = BashOperator(task_id='outside', bash_command=str(datetime.now())) """ warnings = self._check_code(code) - assert len(warnings) == 0, "A task constructed outside the DAG with-block must not be flagged" + assert len(warnings) == 0