Resolve google deprecations in tests (apache#40629)

aws-mwaa · Jul 7, 2024 · be27182 · be27182
1 parent 7360c8b
commit be27182
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 8 deletions.
diff --git a/airflow/providers/google/cloud/sensors/gcs.py b/airflow/providers/google/cloud/sensors/gcs.py
@@ -188,7 +188,15 @@ def ts_function(context):
     try:
         return context["data_interval_end"]
     except KeyError:
-        return context["dag"].following_schedule(context["execution_date"])
+        from airflow.utils import timezone
+
+        data_interval = context["dag"].infer_automated_data_interval(
+            timezone.coerce_datetime(context["execution_date"])
+        )
+        next_info = context["dag"].next_dagrun_info(data_interval, restricted=False)
+        if next_info is None:
+            return None
+        return next_info.data_interval.start
 
 
 class GCSObjectUpdateSensor(BaseSensorOperator):

diff --git a/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py b/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py
@@ -76,7 +76,7 @@ def __init__(
         self.replace_index = replace_index
 
     def get_sql_hook(self) -> PostgresHook:
-        return PostgresHook(schema=self.database, postgres_conn_id=self.postgres_conn_id)
+        return PostgresHook(database=self.database, postgres_conn_id=self.postgres_conn_id)
 
     def execute(self, context: Context) -> None:
         big_query_hook = BigQueryHook(

diff --git a/tests/deprecations_ignore.yml b/tests/deprecations_ignore.yml
@@ -147,12 +147,6 @@
 - tests/providers/google/cloud/operators/test_kubernetes_engine.py::TestGoogleCloudPlatformContainerOperator::test_create_execute_error_body
 - tests/providers/google/cloud/operators/test_life_sciences.py::TestLifeSciencesRunPipelineOperator::test_executes
 - tests/providers/google/cloud/operators/test_life_sciences.py::TestLifeSciencesRunPipelineOperator::test_executes_without_project_id
-- tests/providers/google/cloud/secrets/test_secret_manager.py::TestCloudSecretManagerBackend::test_connections_prefix_none_value
-- tests/providers/google/cloud/secrets/test_secret_manager.py::TestCloudSecretManagerBackend::test_get_conn_uri
-- tests/providers/google/cloud/secrets/test_secret_manager.py::TestCloudSecretManagerBackend::test_get_conn_uri_non_existent_key
-- tests/providers/google/cloud/sensors/test_gcs.py::TestTsFunction::test_should_support_cron
-- tests/providers/google/cloud/sensors/test_gcs.py::TestTsFunction::test_should_support_datetime
-- tests/providers/google/cloud/transfers/test_bigquery_to_postgres.py::TestBigQueryToPostgresOperator::test_execute_good_request_to_bq
 - tests/providers/google/cloud/transfers/test_gcs_to_gcs.py::TestGoogleCloudStorageToCloudStorageOperator::test_copy_files_into_a_folder
 - tests/providers/google/cloud/transfers/test_gcs_to_gcs.py::TestGoogleCloudStorageToCloudStorageOperator::test_execute_last_modified_time
 - tests/providers/google/cloud/transfers/test_gcs_to_gcs.py::TestGoogleCloudStorageToCloudStorageOperator::test_execute_more_than_1_wildcard