diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index aefbe9b9877ae..f417efb23332d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -425,10 +425,15 @@ repos: ^airflow-ctl.*\.py$| ^airflow-core/src/airflow/models/.*\.py$| ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py$| + ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py$| + ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py$| + ^airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py$| + ^airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py$| ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py$| ^airflow-core/tests/unit/cli/commands/test_task_command.py$| ^airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py$| ^airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py$| + ^airflow-core/tests/unit/models/test_deadline.py$| ^airflow-core/tests/unit/models/test_renderedtifields.py$| ^airflow-core/tests/unit/models/test_timestamp.py$| ^airflow-core/tests/unit/utils/test_cli_util.py$| diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py index ca48080c3a0ec..59970da27ea6b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py @@ -137,12 +137,9 @@ def test_should_respond_200_version(self, test_client, accept, session, test_dag # force reserialization test_dag.doc_md = "new doc" force_reserialization(test_dag.dag_id, "dag-folder") - dagcode = ( - session.query(DagCode) - .filter(DagCode.fileloc == test_dag.fileloc) - .order_by(DagCode.id.desc()) - .first() - ) + dagcode = session.scalars( + select(DagCode).where(DagCode.fileloc == test_dag.fileloc).order_by(DagCode.id.desc()) + ).first() assert dagcode.dag_version.version_number == 2 # populate the latest dagcode with a value dag_content2 = "new source code" diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py index d9b5a0287bb4b..93d8c4f63d472 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py @@ -25,7 +25,7 @@ import pytest import time_machine -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from airflow._shared.timezones.timezone import utc, utcnow @@ -275,7 +275,7 @@ def expected_sample_hitl_detail_dict(sample_ti: TaskInstance) -> dict[str, Any]: @pytest.fixture(autouse=True) def cleanup_audit_log(session: Session) -> None: - session.query(Log).delete() + session.execute(delete(Log)) session.commit() diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py index 0a40103c3376f..3e9dd95b3085f 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py @@ -20,6 +20,7 @@ import pendulum import pytest +from sqlalchemy import select from airflow.models.asset import AssetDagRunQueue, AssetEvent, AssetModel from airflow.providers.standard.operators.empty import EmptyOperator @@ -95,7 +96,9 @@ def test_should_set_last_update_only_for_queued_and_hide_flag(self, test_client, assets = { a.uri: a - for a in session.query(AssetModel).filter(AssetModel.uri.in_(["s3://bucket/A", "s3://bucket/B"])) + for a in session.scalars( + select(AssetModel).where(AssetModel.uri.in_(["s3://bucket/A", "s3://bucket/B"])) + ) } # Queue and add an event only for A session.add(AssetDagRunQueue(asset_id=assets["s3://bucket/A"].id, target_dag_id="two_assets_equal")) @@ -145,7 +148,7 @@ def test_last_update_respects_latest_run_filter(self, test_client, dag_maker, se dr = dag_maker.create_dagrun() dag_maker.sync_dagbag_to_db() - asset = session.query(AssetModel).filter(AssetModel.uri == "s3://bucket/F").one() + asset = session.scalars(select(AssetModel).where(AssetModel.uri == "s3://bucket/F")).one() session.add(AssetDagRunQueue(asset_id=asset.id, target_dag_id="filter_run")) # event before latest_run should be ignored ts_base = dr.logical_date or pendulum.now() diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py index 347344dc3598c..080ef093fd9cc 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py @@ -544,7 +544,7 @@ def test_should_return_200_with_resolved_asset_alias_attached_to_the_corrrect_pr self, test_client, session ): resolved_asset = session.scalar( - session.query(AssetModel).filter_by(name="resolved_example_asset_alias") + select(AssetModel).where(AssetModel.name == "resolved_example_asset_alias") ) params = { "dag_id": DAG_ID_RESOLVED_ASSET_ALIAS, diff --git a/airflow-core/tests/unit/models/test_deadline.py b/airflow-core/tests/unit/models/test_deadline.py index d2e575920318a..8521585298c71 100644 --- a/airflow-core/tests/unit/models/test_deadline.py +++ b/airflow-core/tests/unit/models/test_deadline.py @@ -23,6 +23,7 @@ import pytest import time_machine +from sqlalchemy import select from sqlalchemy.exc import SQLAlchemyError from airflow.api_fastapi.core_api.datamodels.dag_run import DAGRunResponse @@ -97,9 +98,9 @@ def dagrun(session, dag_maker): dag_maker.create_dagrun(state=DagRunState.QUEUED, logical_date=DEFAULT_DATE) session.commit() - assert session.query(DagRun).count() == 1 - - return session.query(DagRun).one() + dag_runs = session.scalars(select(DagRun)).all() + assert len(dag_runs) == 1 + return dag_runs[0] @pytest.fixture