ferruzzi commented on code in PR #70714:
URL: https://github.com/apache/airflow/pull/70714#discussion_r3874937299
##########
airflow-core/tests/unit/models/test_deadline.py:
##########
@@ -335,138 +326,68 @@ def
test_handle_miss_persists_executor_callback_routing_data(self, dagrun, sessi
@pytest.mark.db_test
-class TestCalculatedDeadlineDatabaseCalls:
+class TestCalculatedDeadlineReferences:
@staticmethod
def teardown_method():
_clean_db()
@pytest.mark.parametrize(
- ("column", "conditions", "expected_query"),
+ ("reference", "attribute"),
[
pytest.param(
- DagRun.logical_date,
- {"dag_id": DAG_ID},
- "SELECT dag_run.logical_date \nFROM dag_run \nWHERE
dag_run.dag_id = :dag_id_1",
- id="single_condition_logical_date",
- ),
- pytest.param(
- DagRun.queued_at,
- {"dag_id": DAG_ID},
- "SELECT dag_run.queued_at \nFROM dag_run \nWHERE
dag_run.dag_id = :dag_id_1",
- id="single_condition_queued_at",
+ SerializedReferenceModels.DagRunLogicalDateDeadline(),
"logical_date", id="logical_date"
),
+ pytest.param(SerializedReferenceModels.DagRunQueuedAtDeadline(),
"queued_at", id="queued_at"),
pytest.param(
- DagRun.logical_date,
- {"dag_id": DAG_ID, "state": "running"},
- "SELECT dag_run.logical_date \nFROM dag_run \nWHERE
dag_run.dag_id = :dag_id_1 AND dag_run.state = :state_1",
- id="multiple_conditions",
+ ReferenceModels.DagRunLogicalDateDeadline(), "logical_date",
id="legacy_logical_date"
),
+ pytest.param(ReferenceModels.DagRunQueuedAtDeadline(),
"queued_at", id="legacy_queued_at"),
],
)
- @mock.patch("sqlalchemy.orm.Session")
- def test_fetch_from_db_success(self, mock_session, column, conditions,
expected_query):
- """Test successful database queries."""
- mock_session.scalar.return_value = DEFAULT_DATE
-
- result = _fetch_from_db(column, session=mock_session, **conditions)
-
- assert isinstance(result, datetime)
- mock_session.scalar.assert_called_once()
-
- # Check that the correct query was constructed
- call_args = mock_session.scalar.call_args[0][0]
- assert str(call_args) == expected_query
+ def test_dagrun_references_use_supplied_dagrun(self, reference, attribute,
session):
+ """DagRun references use the in-memory DagRun instead of querying it
again."""
+ dagrun = SimpleNamespace(dag_id=DAG_ID, logical_date=DEFAULT_DATE,
queued_at=DEFAULT_DATE)
+ interval = timedelta(hours=1)
- # Verify the actual parameter values
- compiled = call_args.compile()
- for key, value in conditions.items():
- # Note that SQLAlchemy appends the _1 to ensure unique template
field names
- assert compiled.params[f"{key}_1"] == value
+ assert getattr(dagrun, attribute) == DEFAULT_DATE
+ assert (
+ reference.evaluate_with(
+ session=session,
+ interval=interval,
+ dagrun=dagrun,
+ dag_id=DAG_ID,
+ run_id="dagrun_1",
+ unexpected="ignored",
+ )
+ == DEFAULT_DATE + interval
+ )
@pytest.mark.parametrize(
- ("use_valid_conditions", "scalar_side_effect", "expected_error",
"expected_message"),
+ ("reference", "attribute", "message"),
[
pytest.param(
- False,
- mock.DEFAULT, # This will allow the call to pass through
- AttributeError,
- None,
- id="invalid_attribute",
- ),
- pytest.param(
- True,
- SQLAlchemyError("Database connection failed"),
- SQLAlchemyError,
- "Database connection failed",
- id="database_error",
+ SerializedReferenceModels.DagRunLogicalDateDeadline(),
+ "logical_date",
+ "No deadline created for dag_id_1: the Dag run has no logical
date.",
+ id="logical_date",
),
pytest.param(
- True, lambda x: None, ValueError, "No matching record found in
the database", id="no_results"
+ SerializedReferenceModels.DagRunQueuedAtDeadline(),
+ "queued_at",
+ "No deadline created for dag_id_1: the Dag run has no queued
at time.",
+ id="queued_at",
),
],
)
- @mock.patch("sqlalchemy.orm.Session")
- def test_fetch_from_db_error_cases(
- self, mock_session, use_valid_conditions, scalar_side_effect,
expected_error, expected_message
+ def test_dagrun_references_log_when_dagrun_date_is_missing(
+ self, reference, attribute, message, caplog, session
):
- """Test database access error handling."""
- model_reference = DagRun.logical_date
- conditions = {"dag_id": "test_dag"} if use_valid_conditions else
{"non_existent_column": "some_value"}
-
- # Configure mock session
- mock_session.scalar.side_effect = scalar_side_effect
-
- with pytest.raises(expected_error, match=expected_message):
- _fetch_from_db(model_reference, session=mock_session, **conditions)
-
- @pytest.mark.parametrize(
- ("reference", "expected_column"),
- [
- pytest.param(
- SerializedReferenceModels.DagRunLogicalDateDeadline(),
DagRun.logical_date, id="logical_date"
- ),
- pytest.param(
- SerializedReferenceModels.DagRunQueuedAtDeadline(),
DagRun.queued_at, id="queued_at"
- ),
- pytest.param(
- SerializedReferenceModels.FixedDatetimeDeadline(DEFAULT_DATE),
None, id="fixed_deadline"
- ),
- pytest.param(
- SerializedReferenceModels.AverageRuntimeDeadline(max_runs=10,
min_runs=10),
- None,
- id="average_runtime",
- ),
- ],
- )
- def test_deadline_database_integration(self, reference, expected_column,
session):
- """
- Test database integration for all deadline types.
+ caplog.set_level("WARNING", logger=reference.log.name)
+ dagrun = SimpleNamespace(dag_id=DAG_ID, logical_date=DEFAULT_DATE,
queued_at=DEFAULT_DATE)
+ setattr(dagrun, attribute, None)
- Verifies:
- 1. Calculated deadlines call _fetch_from_db with correct column.
- 2. Fixed deadlines do not interact with database.
- 3. Intervals are added to reference times.
- """
- conditions = {"dag_id": DAG_ID, "run_id": "dagrun_1"}
- interval = timedelta(hours=1)
- with
mock.patch("airflow.serialization.definitions.deadline._fetch_from_db") as
mock_fetch:
- mock_fetch.return_value = DEFAULT_DATE
-
- if expected_column is not None:
- result = reference.evaluate_with(session=session,
interval=interval, **conditions)
- mock_fetch.assert_called_once_with(expected_column,
session=session, **conditions)
- elif isinstance(reference,
SerializedReferenceModels.AverageRuntimeDeadline):
- with mock.patch("airflow._shared.timezones.timezone.utcnow")
as mock_utcnow:
- mock_utcnow.return_value = DEFAULT_DATE
- # No DAG runs exist, so it should use 24-hour default
- result = reference.evaluate_with(session=session,
interval=interval, dag_id=DAG_ID)
- mock_fetch.assert_not_called()
- # Should return None when no DAG runs exist
- assert result is None
- else:
- result = reference.evaluate_with(session=session,
interval=interval)
- mock_fetch.assert_not_called()
- assert result == DEFAULT_DATE + interval
+ assert reference.evaluate_with(session=session, interval=timedelta(),
dagrun=dagrun) is None
+ assert caplog.messages == [message]
Review Comment:
Sorry, I know it's frustrating to get follow-up comments like this instead
of getting them all at once. I missed this earlier; we're not supposed to use
caplog.messages anymore. this should be `{"event": ...} in caplog` since we're
using structlog now.
--
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]
