diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 81e24f2e8a3b3..029ad523bfbfe 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -413,6 +413,7 @@ 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/cli/commands/test_rotate_fernet_key_command\.py$| ^task_sdk.*\.py$ pass_filenames: true - id: update-supported-versions diff --git a/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py b/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py index dbd250c2c2d81..142672f184ac0 100644 --- a/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py +++ b/airflow-core/tests/unit/cli/commands/test_rotate_fernet_key_command.py @@ -18,6 +18,7 @@ import pytest from cryptography.fernet import Fernet +from sqlalchemy import select from airflow.cli import cli_parser from airflow.cli.commands import rotate_fernet_key_command @@ -71,7 +72,7 @@ def test_should_rotate_variable(self, session): # Assert correctness using a new fernet key with conf_vars({("core", "fernet_key"): fernet_key2.decode()}): get_fernet.cache_clear() # Clear cached fernet - var1 = session.query(Variable).filter(Variable.key == var1_key).first() + var1 = session.execute(select(Variable).where(Variable.key == var1_key)).scalar_one_or_none() # Unencrypted variable should be unchanged assert Variable.get(key=var1_key) == "value" assert var1._val == "value" @@ -103,7 +104,9 @@ def test_should_rotate_connection(self, session, mock_supervisor_comms): rotate_fernet_key_command.rotate_fernet_key(args) def mock_get_connection(conn_id): - conn = session.query(Connection).filter(Connection.conn_id == conn_id).first() + conn = session.execute( + select(Connection).where(Connection.conn_id == conn_id) + ).scalar_one_or_none() if conn: from airflow.sdk.execution_time.comms import ConnectionResult