From d513e1db778aa79a6356b2c615c231cf9dac515a Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Oct 2022 09:08:54 +0800 Subject: [PATCH 1/7] Fix deprecated max_retries in AWS tests --- airflow/providers/amazon/aws/hooks/athena.py | 4 ++-- airflow/providers/amazon/aws/operators/athena.py | 5 ++++- airflow/providers/amazon/aws/operators/emr.py | 4 +++- airflow/providers/amazon/aws/sensors/emr.py | 6 +++++- tests/providers/amazon/aws/hooks/test_athena.py | 3 ++- tests/providers/amazon/aws/hooks/test_emr_containers.py | 2 +- tests/providers/amazon/aws/operators/test_athena.py | 2 +- tests/providers/amazon/aws/operators/test_emr_containers.py | 2 +- 8 files changed, 19 insertions(+), 9 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/athena.py b/airflow/providers/amazon/aws/hooks/athena.py index 10fa50a3bd94d..c8bac617f6493 100644 --- a/airflow/providers/amazon/aws/hooks/athena.py +++ b/airflow/providers/amazon/aws/hooks/athena.py @@ -207,8 +207,8 @@ def poll_query_status( """ if max_tries: warnings.warn( - f"Method `{self.__class__.__name__}.max_tries` is deprecated and will be removed " - "in a future release. Please use method `max_polling_attempts` instead.", + f"Passing 'max_tries' to {self.__class__.__name__}.poll_query_status is deprecated " + f"and will be removed in a future release. Please use 'max_polling_attempts' instead.", DeprecationWarning, stacklevel=2, ) diff --git a/airflow/providers/amazon/aws/operators/athena.py b/airflow/providers/amazon/aws/operators/athena.py index bed75e224cf53..885f39ec03682 100644 --- a/airflow/providers/amazon/aws/operators/athena.py +++ b/airflow/providers/amazon/aws/operators/athena.py @@ -112,7 +112,10 @@ def execute(self, context: Context) -> str | None: self.client_request_token, self.workgroup, ) - query_status = self.hook.poll_query_status(self.query_execution_id, self.max_polling_attempts) + query_status = self.hook.poll_query_status( + self.query_execution_id, + max_polling_attempts=self.max_polling_attempts, + ) if query_status in AthenaHook.FAILURE_STATES: error_message = self.hook.get_state_change_reason(self.query_execution_id) diff --git a/airflow/providers/amazon/aws/operators/emr.py b/airflow/providers/amazon/aws/operators/emr.py index 659252e10728f..93c8adb8406c7 100644 --- a/airflow/providers/amazon/aws/operators/emr.py +++ b/airflow/providers/amazon/aws/operators/emr.py @@ -279,7 +279,9 @@ def execute(self, context: Context) -> str | None: ) if self.wait_for_completion: query_status = self.hook.poll_query_status( - self.job_id, self.max_polling_attempts, self.poll_interval + self.job_id, + max_polling_attempts=self.max_polling_attempts, + poll_interval=self.poll_interval, ) if query_status in EmrContainerHook.FAILURE_STATES: diff --git a/airflow/providers/amazon/aws/sensors/emr.py b/airflow/providers/amazon/aws/sensors/emr.py index 4759b3d8388a6..97eee6b056e24 100644 --- a/airflow/providers/amazon/aws/sensors/emr.py +++ b/airflow/providers/amazon/aws/sensors/emr.py @@ -285,7 +285,11 @@ def __init__( self.max_retries = max_retries def poke(self, context: Context) -> bool: - state = self.hook.poll_query_status(self.job_id, self.max_retries, self.poll_interval) + state = self.hook.poll_query_status( + self.job_id, + max_polling_attempts=self.max_retries, + poll_interval=self.poll_interval, + ) if state in self.FAILURE_STATES: raise AirflowException('EMR Containers sensor failed') diff --git a/tests/providers/amazon/aws/hooks/test_athena.py b/tests/providers/amazon/aws/hooks/test_athena.py index f8dbd159389a7..f269559d31418 100644 --- a/tests/providers/amazon/aws/hooks/test_athena.py +++ b/tests/providers/amazon/aws/hooks/test_athena.py @@ -164,7 +164,8 @@ def test_hook_poll_query_when_final(self, mock_conn): def test_hook_poll_query_with_timeout(self, mock_conn): mock_conn.return_value.get_query_execution.return_value = MOCK_RUNNING_QUERY_EXECUTION result = self.athena.poll_query_status( - query_execution_id=MOCK_DATA['query_execution_id'], max_tries=1 + query_execution_id=MOCK_DATA['query_execution_id'], + max_polling_attempts=1, ) mock_conn.return_value.get_query_execution.assert_called_once() assert result == 'RUNNING' diff --git a/tests/providers/amazon/aws/hooks/test_emr_containers.py b/tests/providers/amazon/aws/hooks/test_emr_containers.py index 2a14b916c629d..65f35fa5b27a9 100644 --- a/tests/providers/amazon/aws/hooks/test_emr_containers.py +++ b/tests/providers/amazon/aws/hooks/test_emr_containers.py @@ -110,7 +110,7 @@ def test_query_status_polling_with_timeout(self, mock_session): mock_session.return_value = emr_session_mock emr_client_mock.describe_job_run.return_value = JOB2_RUN_DESCRIPTION - query_status = self.emr_containers.poll_query_status(job_id='job123456', max_tries=2) + query_status = self.emr_containers.poll_query_status(job_id='job123456', max_polling_attempts=2) # should poll until max_tries is reached since query is in non-terminal state assert emr_client_mock.describe_job_run.call_count == 2 assert query_status == 'RUNNING' diff --git a/tests/providers/amazon/aws/operators/test_athena.py b/tests/providers/amazon/aws/operators/test_athena.py index ac5d58b10d83d..5fbdebcef4ccc 100644 --- a/tests/providers/amazon/aws/operators/test_athena.py +++ b/tests/providers/amazon/aws/operators/test_athena.py @@ -60,7 +60,7 @@ def setUp(self): output_location='s3://test_s3_bucket/', client_request_token='eac427d0-1c6d-4dfb-96aa-2835d3ac6595', sleep_time=0, - max_tries=3, + max_polling_attempts=3, dag=self.dag, ) diff --git a/tests/providers/amazon/aws/operators/test_emr_containers.py b/tests/providers/amazon/aws/operators/test_emr_containers.py index c5d83a7f5196b..2dc9a660bb812 100644 --- a/tests/providers/amazon/aws/operators/test_emr_containers.py +++ b/tests/providers/amazon/aws/operators/test_emr_containers.py @@ -130,7 +130,7 @@ def test_execute_with_polling_timeout(self, mock_check_query_status): job_driver={}, configuration_overrides={}, poll_interval=0, - max_tries=3, + max_polling_attempts=3, ) with patch('boto3.session.Session', boto3_session_mock): From 09a572b0e22651433e9af4c868f446f436989159 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Oct 2022 09:10:53 +0800 Subject: [PATCH 2/7] Fix connection not found warnings from AWS tests --- .../amazon/aws/hooks/test_batch_client.py | 6 +++ tests/providers/amazon/aws/hooks/test_ec2.py | 3 ++ tests/providers/amazon/aws/hooks/test_eks.py | 3 ++ tests/providers/amazon/aws/hooks/test_s3.py | 6 +++ .../amazon/aws/operators/test_batch.py | 3 ++ .../amazon/aws/operators/test_rds.py | 37 +++++++++++++++ .../amazon/aws/sensors/test_emr_containers.py | 3 ++ .../sensors/test_glue_catalog_partition.py | 3 ++ .../amazon/aws/transfers/test_gcs_to_s3.py | 47 +++++++------------ 9 files changed, 81 insertions(+), 30 deletions(-) diff --git a/tests/providers/amazon/aws/hooks/test_batch_client.py b/tests/providers/amazon/aws/hooks/test_batch_client.py index 8a52096c8da4d..085916f925dd3 100644 --- a/tests/providers/amazon/aws/hooks/test_batch_client.py +++ b/tests/providers/amazon/aws/hooks/test_batch_client.py @@ -53,6 +53,9 @@ def setUp(self, get_client_type_mock): aws_conn_id='airflow_test', region_name=AWS_REGION, ) + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + self.batch_client.get_connection = lambda _: None self.client_mock = get_client_type_mock.return_value assert self.batch_client.client == self.client_mock # setup client property @@ -307,6 +310,9 @@ class TestBatchClientDelays(unittest.TestCase): @mock.patch.dict("os.environ", AWS_SECRET_ACCESS_KEY=AWS_SECRET_ACCESS_KEY) def setUp(self): self.batch_client = BatchClientHook(aws_conn_id='airflow_test', region_name=AWS_REGION) + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + self.batch_client.get_connection = lambda _: None def test_init(self): assert self.batch_client.max_retries == self.batch_client.MAX_RETRIES diff --git a/tests/providers/amazon/aws/hooks/test_ec2.py b/tests/providers/amazon/aws/hooks/test_ec2.py index c0cc75b6b6339..bed2c4c2f4cb1 100644 --- a/tests/providers/amazon/aws/hooks/test_ec2.py +++ b/tests/providers/amazon/aws/hooks/test_ec2.py @@ -31,6 +31,9 @@ def test_init(self): aws_conn_id="aws_conn_test", region_name="region-test", ) + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + ec2_hook.get_connection = lambda _: None assert ec2_hook.aws_conn_id == "aws_conn_test" assert ec2_hook.region_name == "region-test" diff --git a/tests/providers/amazon/aws/hooks/test_eks.py b/tests/providers/amazon/aws/hooks/test_eks.py index 336618b9f418f..3ecb2655806a9 100644 --- a/tests/providers/amazon/aws/hooks/test_eks.py +++ b/tests/providers/amazon/aws/hooks/test_eks.py @@ -1242,6 +1242,9 @@ def test_generate_config_file(self, mock_conn, aws_conn_id, region_name, expecte 'cluster': {'certificateAuthority': {'data': 'test-cert'}, 'endpoint': 'test-endpoint'} } hook = EksHook(aws_conn_id=aws_conn_id, region_name=region_name) + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + hook.get_connection = lambda _: None with hook.generate_config_file( eks_cluster_name='test-cluster', pod_namespace='k8s-namespace' ) as config_file: diff --git a/tests/providers/amazon/aws/hooks/test_s3.py b/tests/providers/amazon/aws/hooks/test_s3.py index 2fb38f77e7e3f..ab5ebfb6d5878 100644 --- a/tests/providers/amazon/aws/hooks/test_s3.py +++ b/tests/providers/amazon/aws/hooks/test_s3.py @@ -46,6 +46,9 @@ def test_check_for_bucket_raises_error_with_invalid_conn_id(self, monkeypatch): monkeypatch.delenv('AWS_ACCESS_KEY_ID', raising=False) monkeypatch.delenv('AWS_SECRET_ACCESS_KEY', raising=False) hook = S3Hook(aws_conn_id="does_not_exist") + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + hook.get_connection = lambda _: None with pytest.raises(NoCredentialsError): hook.check_for_bucket("test-non-existing-bucket") @@ -587,6 +590,9 @@ def test_that_extra_args_not_changed_between_calls(self, s3_bucket): "SSEKMSKeyId": "arn:aws:kms:region:acct-id:key/key-id", } s3_hook = S3Hook(aws_conn_id="s3_test", extra_args=original) + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + s3_hook.get_connection = lambda _: None assert s3_hook.extra_args == original assert s3_hook.extra_args is not original diff --git a/tests/providers/amazon/aws/operators/test_batch.py b/tests/providers/amazon/aws/operators/test_batch.py index f7613f0fc444d..d18281aecf147 100644 --- a/tests/providers/amazon/aws/operators/test_batch.py +++ b/tests/providers/amazon/aws/operators/test_batch.py @@ -66,6 +66,9 @@ def setUp(self, get_client_type_mock): tags={}, ) self.client_mock = self.get_client_type_mock.return_value + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + self.batch.hook.get_connection = lambda _: None assert self.batch.hook.client == self.client_mock # setup client property # don't pause in unit tests diff --git a/tests/providers/amazon/aws/operators/test_rds.py b/tests/providers/amazon/aws/operators/test_rds.py index 729ffd382b40c..b4b3587e502bd 100644 --- a/tests/providers/amazon/aws/operators/test_rds.py +++ b/tests/providers/amazon/aws/operators/test_rds.py @@ -23,6 +23,7 @@ from airflow.exceptions import AirflowException from airflow.models import DAG +from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook from airflow.providers.amazon.aws.hooks.rds import RdsHook from airflow.providers.amazon.aws.operators.rds import ( RdsBaseOperator, @@ -130,6 +131,12 @@ def _create_event_subscription(hook: RdsHook): raise ValueError('AWS not properly mocked') +def _patch_hook_get_connection(hook: AwsGenericHook) -> None: + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + hook.get_connection = lambda _: None + + class TestBaseRdsOperator: dag = None op = None @@ -175,6 +182,7 @@ class TestRdsCreateDbSnapshotOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -192,6 +200,7 @@ def test_create_db_instance_snapshot(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(instance_snapshot_operator.hook) instance_snapshot_operator.execute(None) result = self.hook.conn.describe_db_snapshots(DBSnapshotIdentifier=DB_INSTANCE_SNAPSHOT) @@ -213,6 +222,7 @@ def test_create_db_instance_snapshot_no_wait(self, mock_await_status): dag=self.dag, wait_for_completion=False, ) + _patch_hook_get_connection(instance_snapshot_operator.hook) instance_snapshot_operator.execute(None) result = self.hook.conn.describe_db_snapshots(DBSnapshotIdentifier=DB_INSTANCE_SNAPSHOT) @@ -233,6 +243,7 @@ def test_create_db_cluster_snapshot(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(cluster_snapshot_operator.hook) cluster_snapshot_operator.execute(None) result = self.hook.conn.describe_db_cluster_snapshots(DBClusterSnapshotIdentifier=DB_CLUSTER_SNAPSHOT) @@ -254,6 +265,7 @@ def test_create_db_cluster_snapshot_no_wait(self, mock_no_wait): dag=self.dag, wait_for_completion=False, ) + _patch_hook_get_connection(cluster_snapshot_operator.hook) cluster_snapshot_operator.execute(None) result = self.hook.conn.describe_db_cluster_snapshots(DBClusterSnapshotIdentifier=DB_CLUSTER_SNAPSHOT) @@ -270,6 +282,7 @@ class TestRdsCopyDbSnapshotOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -289,6 +302,7 @@ def test_copy_db_instance_snapshot(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(instance_snapshot_operator.hook) instance_snapshot_operator.execute(None) result = self.hook.conn.describe_db_snapshots(DBSnapshotIdentifier=DB_INSTANCE_SNAPSHOT_COPY) instance_snapshots = result.get("DBSnapshots") @@ -311,6 +325,7 @@ def test_copy_db_instance_snapshot_no_wait(self, mock_await_status): dag=self.dag, wait_for_completion=False, ) + _patch_hook_get_connection(instance_snapshot_operator.hook) instance_snapshot_operator.execute(None) result = self.hook.conn.describe_db_snapshots(DBSnapshotIdentifier=DB_INSTANCE_SNAPSHOT_COPY) instance_snapshots = result.get("DBSnapshots") @@ -332,6 +347,7 @@ def test_copy_db_cluster_snapshot(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(cluster_snapshot_operator.hook) cluster_snapshot_operator.execute(None) result = self.hook.conn.describe_db_cluster_snapshots( DBClusterSnapshotIdentifier=DB_CLUSTER_SNAPSHOT_COPY @@ -355,6 +371,7 @@ def test_copy_db_cluster_snapshot_no_wait(self, mock_await_status): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(cluster_snapshot_operator.hook) cluster_snapshot_operator.execute(None) result = self.hook.conn.describe_db_cluster_snapshots( DBClusterSnapshotIdentifier=DB_CLUSTER_SNAPSHOT_COPY @@ -372,6 +389,7 @@ class TestRdsDeleteDbSnapshotOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -390,6 +408,7 @@ def test_delete_db_instance_snapshot(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(instance_snapshot_operator.hook) instance_snapshot_operator.execute(None) with pytest.raises(self.hook.conn.exceptions.ClientError): @@ -407,6 +426,7 @@ def test_delete_db_cluster_snapshot(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(cluster_snapshot_operator.hook) cluster_snapshot_operator.execute(None) with pytest.raises(self.hook.conn.exceptions.ClientError): @@ -419,6 +439,7 @@ class TestRdsStartExportTaskOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -440,6 +461,7 @@ def test_start_export_task(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(start_export_operator.hook) start_export_operator.execute(None) result = self.hook.conn.describe_export_tasks(ExportTaskIdentifier=EXPORT_TASK_NAME) @@ -466,6 +488,7 @@ def test_start_export_task_no_wait(self, mock_await_status): dag=self.dag, wait_for_completion=False, ) + _patch_hook_get_connection(start_export_operator.hook) start_export_operator.execute(None) result = self.hook.conn.describe_export_tasks(ExportTaskIdentifier=EXPORT_TASK_NAME) @@ -483,6 +506,7 @@ class TestRdsCancelExportTaskOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -501,6 +525,7 @@ def test_cancel_export_task(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(cancel_export_operator.hook) cancel_export_operator.execute(None) result = self.hook.conn.describe_export_tasks(ExportTaskIdentifier=EXPORT_TASK_NAME) @@ -523,6 +548,7 @@ def test_cancel_export_task_no_wait(self, mock_await_status): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(cancel_export_operator.hook) cancel_export_operator.execute(None) result = self.hook.conn.describe_export_tasks(ExportTaskIdentifier=EXPORT_TASK_NAME) @@ -540,6 +566,7 @@ class TestRdsCreateEventSubscriptionOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -559,6 +586,7 @@ def test_create_event_subscription(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(create_subscription_operator.hook) create_subscription_operator.execute(None) result = self.hook.conn.describe_event_subscriptions(SubscriptionName=SUBSCRIPTION_NAME) @@ -582,6 +610,7 @@ def test_create_event_subscription_no_wait(self, mock_await_status): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(create_subscription_operator.hook) create_subscription_operator.execute(None) result = self.hook.conn.describe_event_subscriptions(SubscriptionName=SUBSCRIPTION_NAME) @@ -599,6 +628,7 @@ class TestRdsDeleteEventSubscriptionOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -615,6 +645,7 @@ def test_delete_event_subscription(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(delete_subscription_operator.hook) delete_subscription_operator.execute(None) with pytest.raises(self.hook.conn.exceptions.ClientError): @@ -627,6 +658,7 @@ class TestRdsCreateDbInstanceOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -646,6 +678,7 @@ def test_create_db_instance(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(create_db_instance_operator.hook) create_db_instance_operator.execute(None) result = self.hook.conn.describe_db_instances(DBInstanceIdentifier=DB_INSTANCE_NAME) @@ -669,6 +702,7 @@ def test_create_db_instance_no_wait(self, mock_await_status): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(create_db_instance_operator.hook) create_db_instance_operator.execute(None) result = self.hook.conn.describe_db_instances(DBInstanceIdentifier=DB_INSTANCE_NAME) @@ -686,6 +720,7 @@ class TestRdsDeleteDbInstanceOperator: def setup_class(cls): cls.dag = DAG('test_dag', default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE}) cls.hook = RdsHook(aws_conn_id=AWS_CONN, region_name='us-east-1') + _patch_hook_get_connection(cls.hook) @classmethod def teardown_class(cls): @@ -705,6 +740,7 @@ def test_delete_event_subscription(self): aws_conn_id=AWS_CONN, dag=self.dag, ) + _patch_hook_get_connection(delete_db_instance_operator.hook) delete_db_instance_operator.execute(None) with pytest.raises(self.hook.conn.exceptions.ClientError): @@ -725,6 +761,7 @@ def test_delete_event_subscription_no_wait(self, mock_await_status): dag=self.dag, wait_for_completion=False, ) + _patch_hook_get_connection(delete_db_instance_operator.hook) delete_db_instance_operator.execute(None) with pytest.raises(self.hook.conn.exceptions.ClientError): diff --git a/tests/providers/amazon/aws/sensors/test_emr_containers.py b/tests/providers/amazon/aws/sensors/test_emr_containers.py index 00d2bd8e09db7..1ba326bc9f2e9 100644 --- a/tests/providers/amazon/aws/sensors/test_emr_containers.py +++ b/tests/providers/amazon/aws/sensors/test_emr_containers.py @@ -37,6 +37,9 @@ def setUp(self): max_retries=1, aws_conn_id='aws_default', ) + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + self.sensor.hook.get_connection = lambda _: None @mock.patch.object(EmrContainerHook, 'check_query_status', side_effect=("PENDING",)) def test_poke_pending(self, mock_check_query_status): diff --git a/tests/providers/amazon/aws/sensors/test_glue_catalog_partition.py b/tests/providers/amazon/aws/sensors/test_glue_catalog_partition.py index 8592574905134..9f0db91b1b53b 100644 --- a/tests/providers/amazon/aws/sensors/test_glue_catalog_partition.py +++ b/tests/providers/amazon/aws/sensors/test_glue_catalog_partition.py @@ -79,6 +79,9 @@ def test_poke_nondefault_args(self, mock_check_for_partition): poke_interval=poke_interval, timeout=timeout, ) + # We're mocking all actual AWS calls and don't need a connection. This + # avoids an Airflow warning about connection cannot be found. + op.get_hook().get_connection = lambda _: None op.poke({}) assert op.hook.region_name == region_name diff --git a/tests/providers/amazon/aws/transfers/test_gcs_to_s3.py b/tests/providers/amazon/aws/transfers/test_gcs_to_s3.py index 5ddc70628435f..bace72ddf51ac 100644 --- a/tests/providers/amazon/aws/transfers/test_gcs_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_gcs_to_s3.py @@ -38,6 +38,16 @@ S3_ACL_POLICY = "private-read" +def _create_test_bucket(): + hook = S3Hook(aws_conn_id='airflow_gcs_test') + # We're mocking all actual AWS calls and don't need a connection. + # This avoids an Airflow warning about connection cannot be found. + hook.get_connection = lambda _: None + bucket = hook.get_bucket('bucket') + bucket.create() + return hook, bucket + + class TestGCSToS3Operator(unittest.TestCase): # Test1: incremental behaviour (just some files missing) @@ -58,10 +68,7 @@ def test_execute_incremental(self, mock_hook): dest_s3_key=S3_BUCKET, replace=False, ) - # create dest bucket - hook = S3Hook(aws_conn_id='airflow_gcs_test') - bucket = hook.get_bucket('bucket') - bucket.create() + hook, bucket = _create_test_bucket() bucket.put_object(Key=MOCK_FILES[0], Body=b'testing') # we expect all except first file in MOCK_FILES to be uploaded @@ -88,10 +95,7 @@ def test_execute_without_replace(self, mock_hook): dest_s3_key=S3_BUCKET, replace=False, ) - # create dest bucket with all the files - hook = S3Hook(aws_conn_id='airflow_gcs_test') - bucket = hook.get_bucket('bucket') - bucket.create() + hook, bucket = _create_test_bucket() for mock_file in MOCK_FILES: bucket.put_object(Key=mock_file, Body=b'testing') @@ -119,10 +123,7 @@ def test_execute(self, mock_hook): dest_s3_key=S3_BUCKET, replace=False, ) - # create dest bucket without files - hook = S3Hook(aws_conn_id='airflow_gcs_test') - bucket = hook.get_bucket('bucket') - bucket.create() + hook, _ = _create_test_bucket() # we expect all MOCK_FILES to be uploaded # and all MOCK_FILES to be present at the S3 bucket @@ -148,10 +149,7 @@ def test_execute_with_replace(self, mock_hook): dest_s3_key=S3_BUCKET, replace=True, ) - # create dest bucket with all the files - hook = S3Hook(aws_conn_id='airflow_gcs_test') - bucket = hook.get_bucket('bucket') - bucket.create() + hook, bucket = _create_test_bucket() for mock_file in MOCK_FILES: bucket.put_object(Key=mock_file, Body=b'testing') @@ -179,10 +177,7 @@ def test_execute_incremental_with_replace(self, mock_hook): dest_s3_key=S3_BUCKET, replace=True, ) - # create dest bucket with just two files (the first two files in MOCK_FILES) - hook = S3Hook(aws_conn_id='airflow_gcs_test') - bucket = hook.get_bucket('bucket') - bucket.create() + hook, bucket = _create_test_bucket() for mock_file in MOCK_FILES[:2]: bucket.put_object(Key=mock_file, Body=b'testing') @@ -261,12 +256,7 @@ def test_execute_with_s3_acl_policy(self, mock_load_file, mock_gcs_hook): replace=False, s3_acl_policy=S3_ACL_POLICY, ) - - # Create dest bucket without files - hook = S3Hook(aws_conn_id='airflow_gcs_test') - bucket = hook.get_bucket('bucket') - bucket.create() - + _create_test_bucket() operator.execute(None) # Make sure the acl_policy parameter is passed to the upload method @@ -291,10 +281,7 @@ def test_execute_without_keep_director_structure(self, mock_hook): replace=False, keep_directory_structure=False, ) - # create dest bucket with all the files - hook = S3Hook(aws_conn_id='airflow_gcs_test') - bucket = hook.get_bucket('bucket') - bucket.create() + hook, _ = _create_test_bucket() # we expect all except first file in MOCK_FILES to be uploaded # and all the MOCK_FILES to be present at the S3 bucket From b40b14a0af37dd2a6e99a159bfa7974a5bc3f4e9 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Oct 2022 09:23:52 +0800 Subject: [PATCH 3/7] Fix deprecation for SQLAlchemy's URL class --- airflow/providers/amazon/aws/hooks/redshift_sql.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/hooks/redshift_sql.py b/airflow/providers/amazon/aws/hooks/redshift_sql.py index 608899a914dc1..c7860928854f4 100644 --- a/airflow/providers/amazon/aws/hooks/redshift_sql.py +++ b/airflow/providers/amazon/aws/hooks/redshift_sql.py @@ -82,7 +82,10 @@ def get_uri(self) -> str: if 'user' in conn_params: conn_params['username'] = conn_params.pop('user') - return str(URL(drivername='redshift+redshift_connector', **conn_params)) + # Compatibility: The 'create' factory method was added in SQLAlchemy 1.4 + # to replace calling the default URL constructor directly. + create_url = getattr(URL, "create", URL) + return str(create_url(drivername='redshift+redshift_connector', **conn_params)) def get_sqlalchemy_engine(self, engine_kwargs=None): """Overrides DbApiHook get_sqlalchemy_engine to pass redshift_connector specific kwargs""" From 51d1e7b65ea2214c6afeabf8e8ceaba8313b716e Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Oct 2022 09:34:22 +0800 Subject: [PATCH 4/7] Fix deprecated Watchtower arguments We are already on watchtower>=2.0 so the new arguments should always be available. --- .../amazon/aws/log/cloudwatch_task_handler.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py b/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py index 9d4b7fc247013..54d8a0df522a2 100644 --- a/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py +++ b/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py @@ -63,8 +63,8 @@ def _render_filename(self, ti, try_number): def set_context(self, ti): super().set_context(ti) self.handler = watchtower.CloudWatchLogHandler( - log_group=self.log_group, - stream_name=self._render_filename(ti, ti.try_number), + log_group_name=self.log_group, + log_stream_name=self._render_filename(ti, ti.try_number), boto3_client=self.hook.get_conn(), ) @@ -98,12 +98,11 @@ def get_cloudwatch_logs(self, stream_name: str) -> str: :return: string of all logs from the given log stream """ try: - events = list( - self.hook.get_log_events( - log_group=self.log_group, log_stream_name=stream_name, start_from_head=True - ) + events = self.hook.get_log_events( + log_group=self.log_group, + log_stream_name=stream_name, + start_from_head=True, ) - return '\n'.join(self._event_to_str(event) for event in events) except Exception: msg = f'Could not read remote logs from log_group: {self.log_group} log_stream: {stream_name}.' From a78f7a87af67ee899b18a47a0e5a04bcb082759c Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Oct 2022 09:40:07 +0800 Subject: [PATCH 5/7] Fix deprecated load_test_config usage --- tests/providers/amazon/aws/operators/test_emr_containers.py | 6 +++--- .../providers/amazon/aws/transfers/test_google_api_to_s3.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/providers/amazon/aws/operators/test_emr_containers.py b/tests/providers/amazon/aws/operators/test_emr_containers.py index 2dc9a660bb812..866b898c47765 100644 --- a/tests/providers/amazon/aws/operators/test_emr_containers.py +++ b/tests/providers/amazon/aws/operators/test_emr_containers.py @@ -22,7 +22,7 @@ import pytest -from airflow import configuration +from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook from airflow.providers.amazon.aws.operators.emr import EmrContainerOperator, EmrEksCreateClusterOperator @@ -41,7 +41,7 @@ class TestEmrContainerOperator(unittest.TestCase): @mock.patch('airflow.providers.amazon.aws.hooks.emr.EmrContainerHook') def setUp(self, emr_hook_mock): - configuration.load_test_config() + conf.load_test_config() self.emr_hook_mock = emr_hook_mock self.emr_container = EmrContainerOperator( @@ -145,7 +145,7 @@ def test_execute_with_polling_timeout(self, mock_check_query_status): class TestEmrEksCreateClusterOperator(unittest.TestCase): @mock.patch('airflow.providers.amazon.aws.hooks.emr.EmrContainerHook') def setUp(self, emr_hook_mock): - configuration.load_test_config() + conf.load_test_config() self.emr_hook_mock = emr_hook_mock self.emr_container = EmrEksCreateClusterOperator( diff --git a/tests/providers/amazon/aws/transfers/test_google_api_to_s3.py b/tests/providers/amazon/aws/transfers/test_google_api_to_s3.py index 07cf9bcc0d09a..4c0abd86dd694 100644 --- a/tests/providers/amazon/aws/transfers/test_google_api_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_google_api_to_s3.py @@ -23,7 +23,7 @@ import pytest from airflow import models -from airflow.configuration import load_test_config +from airflow.configuration import conf from airflow.models.xcom import MAX_XCOM_SIZE from airflow.providers.amazon.aws.transfers.google_api_to_s3 import GoogleApiToS3Operator from airflow.utils import db @@ -31,7 +31,7 @@ class TestGoogleApiToS3(unittest.TestCase): def setUp(self): - load_test_config() + conf.load_test_config() db.merge_conn( models.Connection( From 67ce94d1edbcba214490699119fc3eafabacee06 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Oct 2022 11:20:59 +0800 Subject: [PATCH 6/7] Fix deprecation warnings in secrets manager tests --- .../amazon/aws/secrets/secrets_manager.py | 14 ++-- .../aws/secrets/test_secrets_manager.py | 67 ++++++++++++------- .../aws/secrets/test_systems_manager.py | 12 ++-- 3 files changed, 54 insertions(+), 39 deletions(-) diff --git a/airflow/providers/amazon/aws/secrets/secrets_manager.py b/airflow/providers/amazon/aws/secrets/secrets_manager.py index 99192504b8454..3d53dbeaf2570 100644 --- a/airflow/providers/amazon/aws/secrets/secrets_manager.py +++ b/airflow/providers/amazon/aws/secrets/secrets_manager.py @@ -330,19 +330,19 @@ def _remove_escaping_in_secret_dict(self, secret: dict[str, Any], conn_id: str) if warn_user: msg = ( - "When ``full_url_mode=True``, URL-encoding secret values is deprecated. In future versions, " - f" this value will not be un-escaped. For the conn_id {conn_id!r}, please remove the" - " URL-encoding." - "\n\nThis warning was raised because the SecretsManagerBackend detected that this connection" - " was URL-encoded." + "When full_url_mode=False, URL-encoding secret values is deprecated. In future versions, " + f"this value will not be un-escaped. For the conn_id {conn_id!r}, please remove the " + "URL-encoding.\n\n" + "This warning was raised because the SecretsManagerBackend detected that this " + "connection was URL-encoded." ) if idempotent: msg = f" Once the values for conn_id {conn_id!r} are decoded, this warning will go away." if not idempotent: msg += ( " In addition to decoding the values for your connection, you must also set" - " ``secret_values_are_urlencoded=False`` for your config variable" - " ``secrets.backend_kwargs`` because this connection's URL encoding is not idempotent." + " secret_values_are_urlencoded=False for your config variable" + " secrets.backend_kwargs because this connection's URL encoding is not idempotent." " For more information, see:" " https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/secrets-backends" "/aws-secrets-manager.html#url-encoding-of-secrets-when-full-url-mode-is-false" diff --git a/tests/providers/amazon/aws/secrets/test_secrets_manager.py b/tests/providers/amazon/aws/secrets/test_secrets_manager.py index c3cb890db15a9..9b30f5aa89f02 100644 --- a/tests/providers/amazon/aws/secrets/test_secrets_manager.py +++ b/tests/providers/amazon/aws/secrets/test_secrets_manager.py @@ -17,14 +17,15 @@ from __future__ import annotations import json -from unittest import TestCase, mock +from unittest import mock +import pytest from moto import mock_secretsmanager from airflow.providers.amazon.aws.secrets.secrets_manager import SecretsManagerBackend -class TestSecretsManagerBackend(TestCase): +class TestSecretsManagerBackend: @mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend.get_conn_value") def test_aws_secrets_manager_get_connection(self, mock_get_value): mock_get_value.return_value = "scheme://user:pass@host:100" @@ -32,7 +33,7 @@ def test_aws_secrets_manager_get_connection(self, mock_get_value): assert conn.host == 'host' @mock_secretsmanager - def test_get_conn_uri_full_url_mode(self): + def test_get_conn_value_full_url_mode(self): secret_id = 'airflow/connections/test_postgres' create_param = { 'Name': secret_id, @@ -47,11 +48,18 @@ def test_get_conn_uri_full_url_mode(self): secrets_manager_backend.client.create_secret(**create_param) secrets_manager_backend.client.put_secret_value(**param) - returned_uri = secrets_manager_backend.get_conn_uri(conn_id="test_postgres") + returned_uri = secrets_manager_backend.get_conn_value(conn_id="test_postgres") assert 'postgresql://airflow:airflow@host:5432/airflow' == returned_uri + @pytest.mark.parametrize( + "full_url_mode, login, host", + [ + (False, "is url encoded", "not%20idempotent"), + (True, "is%20url%20encoded", "not%2520idempotent"), + ], + ) @mock_secretsmanager - def test_get_connection_broken_field_mode_url_encoding(self): + def test_get_connection_broken_field_mode_url_encoding(self, full_url_mode, login, host): secret_id = 'airflow/connections/test_postgres' create_param = { 'Name': secret_id, @@ -62,7 +70,7 @@ def test_get_connection_broken_field_mode_url_encoding(self): 'SecretString': json.dumps( { 'conn_type': 'postgresql', - 'user': 'is%20url%20encoded', + 'login': 'is%20url%20encoded', 'password': 'not url encoded', 'host': 'not%2520idempotent', 'extra': json.dumps({'foo': 'bar'}), @@ -70,23 +78,20 @@ def test_get_connection_broken_field_mode_url_encoding(self): ), } - secrets_manager_backend = SecretsManagerBackend(full_url_mode=False) + secrets_manager_backend = SecretsManagerBackend(full_url_mode=full_url_mode) secrets_manager_backend.client.create_secret(**create_param) secrets_manager_backend.client.put_secret_value(**param) - conn = secrets_manager_backend.get_connection(conn_id='test_postgres') - assert conn.login == 'is url encoded' - assert conn.password == 'not url encoded' - assert conn.host == 'not%20idempotent' - assert conn.conn_id == 'test_postgres' - - # Remove URL encoding - secrets_manager_backend.are_secret_values_urlencoded = False + if full_url_mode: + conn = secrets_manager_backend.get_connection(conn_id='test_postgres') + else: + warning_match = r"When full_url_mode=False, URL-encoding secret values is deprecated\..+" + with pytest.warns(DeprecationWarning, match=warning_match): + conn = secrets_manager_backend.get_connection(conn_id='test_postgres') - conn = secrets_manager_backend.get_connection(conn_id='test_postgres') - assert conn.login == 'is%20url%20encoded' + assert conn.login == login assert conn.password == 'not url encoded' - assert conn.host == 'not%2520idempotent' + assert conn.host == host assert conn.conn_id == 'test_postgres' assert conn.extra_dejson['foo'] == 'bar' @@ -118,7 +123,7 @@ def test_get_connection_broken_field_mode_extra_allows_nested_json(self): assert conn.extra_dejson['foo'] == 'bar' @mock_secretsmanager - def test_get_conn_uri_broken_field_mode(self): + def test_get_conn_value_broken_field_mode(self): secret_id = 'airflow/connections/test_postgres' create_param = { 'Name': secret_id, @@ -134,11 +139,16 @@ def test_get_conn_uri_broken_field_mode(self): secrets_manager_backend.client.create_secret(**create_param) secrets_manager_backend.client.put_secret_value(**param) - returned_uri = secrets_manager_backend.get_conn_uri(conn_id="test_postgres") + warning_match = ( + r"In future versions, `SecretsManagerBackend\.get_conn_value` will return a JSON string when " + r"full_url_mode is False, not a URI\." + ) + with pytest.warns(DeprecationWarning, match=warning_match): + returned_uri = secrets_manager_backend.get_conn_value(conn_id="test_postgres") assert 'postgresql://airflow:airflow@host:5432/airflow' == returned_uri @mock_secretsmanager - def test_get_conn_uri_broken_field_mode_extra_words_added(self): + def test_get_conn_value_broken_field_mode_extra_words_added(self): secret_id = 'airflow/connections/test_postgres' create_param = { 'Name': secret_id, @@ -156,7 +166,12 @@ def test_get_conn_uri_broken_field_mode_extra_words_added(self): secrets_manager_backend.client.create_secret(**create_param) secrets_manager_backend.client.put_secret_value(**param) - returned_uri = secrets_manager_backend.get_conn_uri(conn_id="test_postgres") + warning_match = ( + r"In future versions, `SecretsManagerBackend\.get_conn_value` will return a JSON string when " + r"full_url_mode is False, not a URI\." + ) + with pytest.warns(DeprecationWarning, match=warning_match): + returned_uri = secrets_manager_backend.get_conn_value(conn_id="test_postgres") assert 'postgresql://airflow:airflow@host:5432/airflow' == returned_uri @mock_secretsmanager @@ -170,7 +185,7 @@ def test_format_uri_with_extra(self): assert conn_string_with_extra == 'CS?key1=value1&key2=value2' @mock_secretsmanager - def test_get_conn_uri_non_existent_key(self): + def test_get_conn_value_non_existent_key(self): """ Test that if the key with connection ID is not present, SecretsManagerBackend.get_connection should return None @@ -191,7 +206,7 @@ def test_get_conn_uri_non_existent_key(self): secrets_manager_backend.client.create_secret(**create_param) secrets_manager_backend.client.put_secret_value(**param) - assert secrets_manager_backend.get_conn_uri(conn_id=conn_id) is None + assert secrets_manager_backend.get_conn_value(conn_id=conn_id) is None assert secrets_manager_backend.get_connection(conn_id=conn_id) is None @mock_secretsmanager @@ -233,14 +248,14 @@ def test_get_variable_non_existent_key(self): def test_connection_prefix_none_value(self, mock_get_secret): """ Test that if Variable key is not present in AWS Secrets Manager, - SecretsManagerBackend.get_conn_uri should return None, + SecretsManagerBackend.get_conn_value should return None, SecretsManagerBackend._get_secret should not be called """ kwargs = {'connections_prefix': None} secrets_manager_backend = SecretsManagerBackend(**kwargs) - assert secrets_manager_backend.get_conn_uri("test_mysql") is None + assert secrets_manager_backend.get_conn_value("test_mysql") is None @mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend._get_secret") def test_variable_prefix_none_value(self, mock_get_secret): diff --git a/tests/providers/amazon/aws/secrets/test_systems_manager.py b/tests/providers/amazon/aws/secrets/test_systems_manager.py index bdcc36337f180..92832b1e2cb1d 100644 --- a/tests/providers/amazon/aws/secrets/test_systems_manager.py +++ b/tests/providers/amazon/aws/secrets/test_systems_manager.py @@ -36,7 +36,7 @@ def test_aws_ssm_get_connection(self, mock_get_value): assert conn.host == 'host' @mock_ssm - def test_get_conn_uri(self): + def test_get_conn_value(self): param = { 'Name': '/airflow/connections/test_postgres', 'Type': 'String', @@ -46,11 +46,11 @@ def test_get_conn_uri(self): ssm_backend = SystemsManagerParameterStoreBackend() ssm_backend.client.put_parameter(**param) - returned_uri = ssm_backend.get_conn_uri(conn_id="test_postgres") + returned_uri = ssm_backend.get_conn_value(conn_id="test_postgres") assert 'postgresql://airflow:airflow@host:5432/airflow' == returned_uri @mock_ssm - def test_get_conn_uri_non_existent_key(self): + def test_get_conn_value_non_existent_key(self): """ Test that if the key with connection ID is not present in SSM, SystemsManagerParameterStoreBackend.get_connection should return None @@ -65,7 +65,7 @@ def test_get_conn_uri_non_existent_key(self): ssm_backend = SystemsManagerParameterStoreBackend() ssm_backend.client.put_parameter(**param) - assert ssm_backend.get_conn_uri(conn_id=conn_id) is None + assert ssm_backend.get_conn_value(conn_id=conn_id) is None assert ssm_backend.get_connection(conn_id=conn_id) is None @mock_ssm @@ -158,14 +158,14 @@ def test_passing_client_kwargs(self, mock_session_factory): def test_connection_prefix_none_value(self, mock_get_secret): """ Test that if Variable key is not present in SSM, - SystemsManagerParameterStoreBackend.get_conn_uri should return None, + SystemsManagerParameterStoreBackend.get_conn_value should return None, SystemsManagerParameterStoreBackend._get_secret should not be called """ kwargs = {'connections_prefix': None} ssm_backend = SystemsManagerParameterStoreBackend(**kwargs) - assert ssm_backend.get_conn_uri("test_mysql") is None + assert ssm_backend.get_conn_value("test_mysql") is None mock_get_secret.assert_not_called() @mock.patch( From 8c13cb8eda58f71f65e52ecf9fc1a6839c296dd3 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 4 Oct 2022 13:30:27 +0800 Subject: [PATCH 7/7] Expect warnings in AWS connection wrapper tests --- tests/providers/amazon/aws/utils/test_connection_wrapper.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/providers/amazon/aws/utils/test_connection_wrapper.py b/tests/providers/amazon/aws/utils/test_connection_wrapper.py index 3540f86e6154e..366246eb9d679 100644 --- a/tests/providers/amazon/aws/utils/test_connection_wrapper.py +++ b/tests/providers/amazon/aws/utils/test_connection_wrapper.py @@ -44,7 +44,7 @@ def test_values_from_connection(self, extra): password="mock-password", extra=extra, # AwsBaseHook never use this attributes from airflow.models.Connection - host="mock-host", + host=None, schema="mock-schema", port=42, ) @@ -141,7 +141,8 @@ def test_get_credentials_from_session_kwargs( } mock_conn = mock_connection_factory(login=None, password=None, extra=mock_conn_extra) - wrap_conn = AwsConnectionWrapper(conn=mock_conn) + with pytest.warns(DeprecationWarning, match=r"'session_kwargs' in extra config is deprecated"): + wrap_conn = AwsConnectionWrapper(conn=mock_conn) assert wrap_conn.aws_access_key_id == aws_access_key_id assert wrap_conn.aws_secret_access_key == aws_secret_access_key assert wrap_conn.aws_session_token == aws_session_token