diff --git a/airflow/contrib/hooks/gcs_hook.py b/airflow/contrib/hooks/gcs_hook.py index 7ccaa4f028db3..ee57875530373 100644 --- a/airflow/contrib/hooks/gcs_hook.py +++ b/airflow/contrib/hooks/gcs_hook.py @@ -22,9 +22,23 @@ import warnings # pylint: disable=unused-import -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook, _parse_gcs_url # noqa +from airflow.gcp.hooks.gcs import GcsHook, _parse_gcs_url # noqa warnings.warn( "This module is deprecated. Please use `airflow.gcp.hooks.gcs`.", DeprecationWarning, stacklevel=2 ) + + +class GoogleCloudStorageHook(GcsHook): + """ + This class is deprecated. Please use `airflow.gcp.hooks.gcs.GcsHook`. + """ + + def __init__(self, *args, **kwargs): + warnings.warn( + "This class is deprecated. Please use `airflow.gcp.hooks.gcs.GcsHook`.", + DeprecationWarning, stacklevel=2 + ) + + super().__init__(*args, **kwargs) diff --git a/airflow/contrib/operators/gcs_to_gdrive_operator.py b/airflow/contrib/operators/gcs_to_gdrive_operator.py index 2ff7a6a487469..f48db9c5dba0c 100644 --- a/airflow/contrib/operators/gcs_to_gdrive_operator.py +++ b/airflow/contrib/operators/gcs_to_gdrive_operator.py @@ -24,7 +24,7 @@ from airflow.contrib.hooks.gdrive_hook import GoogleDriveHook from airflow.exceptions import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -93,12 +93,12 @@ def __init__( self.move_object = move_object self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to - self.gcs_hook = None # type: Optional[GoogleCloudStorageHook] + self.gcs_hook = None # type: Optional[GcsHook] self.gdrive_hook = None # type: Optional[GoogleDriveHook] def execute(self, context): - self.gcs_hook = GoogleCloudStorageHook( + self.gcs_hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) self.gdrive_hook = GoogleDriveHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) diff --git a/airflow/contrib/operators/s3_to_gcs_operator.py b/airflow/contrib/operators/s3_to_gcs_operator.py index 0c86e481da95d..a99f04f63abfd 100644 --- a/airflow/contrib/operators/s3_to_gcs_operator.py +++ b/airflow/contrib/operators/s3_to_gcs_operator.py @@ -21,7 +21,7 @@ from airflow.contrib.operators.s3_list_operator import S3ListOperator from airflow.exceptions import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook, _parse_gcs_url +from airflow.gcp.hooks.gcs import GcsHook, _parse_gcs_url from airflow.hooks.S3_hook import S3Hook from airflow.utils.decorators import apply_defaults @@ -140,7 +140,7 @@ def execute(self, context): # use the super method to list all the files in an S3 bucket/key files = super().execute(context) - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) diff --git a/airflow/gcp/hooks/gcs.py b/airflow/gcp/hooks/gcs.py index efb206069f58d..5960af304b01a 100644 --- a/airflow/gcp/hooks/gcs.py +++ b/airflow/gcp/hooks/gcs.py @@ -36,7 +36,7 @@ from airflow.version import version -class GoogleCloudStorageHook(GoogleCloudBaseHook): +class GcsHook(GoogleCloudBaseHook): """ Interact with Google Cloud Storage. This hook uses the Google Cloud Platform connection. diff --git a/airflow/gcp/operators/bigquery.py b/airflow/gcp/operators/bigquery.py index a3270608582af..d01d138b9ac28 100644 --- a/airflow/gcp/operators/bigquery.py +++ b/airflow/gcp/operators/bigquery.py @@ -28,7 +28,7 @@ from airflow.exceptions import AirflowException from airflow.gcp.hooks.bigquery import BigQueryHook -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook, _parse_gcs_url +from airflow.gcp.hooks.gcs import GcsHook, _parse_gcs_url from airflow.models.baseoperator import BaseOperator, BaseOperatorLink from airflow.models.taskinstance import TaskInstance from airflow.operators.check_operator import CheckOperator, IntervalCheckOperator, ValueCheckOperator @@ -734,7 +734,7 @@ def execute(self, context): gcs_bucket, gcs_object = _parse_gcs_url(self.gcs_schema_object) - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( google_cloud_storage_conn_id=self.google_cloud_storage_conn_id, delegate_to=self.delegate_to) schema_fields = json.loads(gcs_hook.download( @@ -903,7 +903,7 @@ def execute(self, context): location=self.location) if not self.schema_fields and self.schema_object and self.source_format != 'DATASTORE_BACKUP': - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( google_cloud_storage_conn_id=self.google_cloud_storage_conn_id, delegate_to=self.delegate_to) schema_fields = json.loads(gcs_hook.download( diff --git a/airflow/gcp/operators/dataflow.py b/airflow/gcp/operators/dataflow.py index 534ea503e1246..5a2652f59cb56 100644 --- a/airflow/gcp/operators/dataflow.py +++ b/airflow/gcp/operators/dataflow.py @@ -29,7 +29,7 @@ from typing import List, Optional from airflow.gcp.hooks.dataflow import DataFlowHook -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults from airflow.version import version @@ -422,13 +422,13 @@ def execute(self, context): class GoogleCloudBucketHelper: - """GoogleCloudStorageHook helper class to download GCS object.""" + """GcsHook helper class to download GCS object.""" GCS_PREFIX_LENGTH = 5 def __init__(self, gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None) -> None: - self._gcs_hook = GoogleCloudStorageHook(gcp_conn_id, delegate_to) + self._gcs_hook = GcsHook(gcp_conn_id, delegate_to) def google_cloud_to_local(self, file_name: str) -> str: """ diff --git a/airflow/gcp/operators/dataproc.py b/airflow/gcp/operators/dataproc.py index 91b84c95d2b67..7f4f5f610706a 100644 --- a/airflow/gcp/operators/dataproc.py +++ b/airflow/gcp/operators/dataproc.py @@ -32,7 +32,7 @@ from airflow.exceptions import AirflowException from airflow.gcp.hooks.dataproc import DataProcHook -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils import timezone from airflow.utils.decorators import apply_defaults @@ -1105,7 +1105,7 @@ def _upload_file_temp(self, bucket, local_file): self.log.info("Uploading %s to %s", local_file, temp_filename) - GoogleCloudStorageHook( + GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id ).upload( bucket_name=bucket, diff --git a/airflow/gcp/operators/datastore.py b/airflow/gcp/operators/datastore.py index 313e0b087e4f3..8e07793ffc3ce 100644 --- a/airflow/gcp/operators/datastore.py +++ b/airflow/gcp/operators/datastore.py @@ -24,7 +24,7 @@ from airflow.exceptions import AirflowException from airflow.gcp.hooks.datastore import DatastoreHook -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -94,7 +94,7 @@ def execute(self, context): self.log.info('Exporting data to Cloud Storage bucket %s', self.bucket) if self.overwrite_existing and self.namespace: - gcs_hook = GoogleCloudStorageHook(self.cloud_storage_conn_id) + gcs_hook = GcsHook(self.cloud_storage_conn_id) objects = gcs_hook.list(self.bucket, prefix=self.namespace) for obj in objects: gcs_hook.delete(self.bucket, obj) diff --git a/airflow/gcp/operators/gcs.py b/airflow/gcp/operators/gcs.py index efbe264dfe628..9112a80bad81d 100644 --- a/airflow/gcp/operators/gcs.py +++ b/airflow/gcp/operators/gcs.py @@ -26,7 +26,7 @@ from typing import Dict, Iterable, List, Optional, Union from airflow import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.models.xcom import MAX_XCOM_SIZE from airflow.utils.decorators import apply_defaults @@ -130,7 +130,7 @@ def __init__(self, self.delegate_to = delegate_to def execute(self, context): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) @@ -211,7 +211,7 @@ def __init__(self, def execute(self, context): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) @@ -300,7 +300,7 @@ def __init__(self, def execute(self, context): self.log.info('Executing download: %s, %s, %s', self.bucket, self.object, self.filename) - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) @@ -373,7 +373,7 @@ def __init__(self, super().__init__(*args, **kwargs) def execute(self, context): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) @@ -448,7 +448,7 @@ def __init__( self.gcp_conn_id = gcp_conn_id def execute(self, context): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id ) hook.insert_bucket_acl(bucket_name=self.bucket, entity=self.entity, role=self.role, @@ -519,7 +519,7 @@ def __init__(self, self.gcp_conn_id = gcp_conn_id def execute(self, context): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id ) hook.insert_object_acl(bucket_name=self.bucket, @@ -580,7 +580,7 @@ def __init__( self.output_encoding = sys.getdefaultencoding() def execute(self, context: Dict): - hook = GoogleCloudStorageHook(gcp_conn_id=self.gcp_conn_id) + hook = GcsHook(gcp_conn_id=self.gcp_conn_id) with NamedTemporaryFile() as source_file, NamedTemporaryFile() as destination_file: self.log.info("Downloading file from %s", self.source_bucket) diff --git a/airflow/gcp/operators/text_to_speech.py b/airflow/gcp/operators/text_to_speech.py index 0a8b93021774f..16003c247bf0f 100644 --- a/airflow/gcp/operators/text_to_speech.py +++ b/airflow/gcp/operators/text_to_speech.py @@ -27,7 +27,7 @@ from google.cloud.texttospeech_v1.types import AudioConfig, SynthesisInput, VoiceSelectionParams from airflow import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.gcp.hooks.text_to_speech import GCPTextToSpeechHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -129,7 +129,7 @@ def execute(self, context): ) with NamedTemporaryFile() as temp_file: temp_file.write(result.audio_content) - cloud_storage_hook = GoogleCloudStorageHook(google_cloud_storage_conn_id=self.gcp_conn_id) + cloud_storage_hook = GcsHook(google_cloud_storage_conn_id=self.gcp_conn_id) cloud_storage_hook.upload( bucket_name=self.target_bucket_name, object_name=self.target_filename, filename=temp_file.name ) diff --git a/airflow/gcp/sensors/gcs.py b/airflow/gcp/sensors/gcs.py index 3d16ee7dec322..06317ce533684 100644 --- a/airflow/gcp/sensors/gcs.py +++ b/airflow/gcp/sensors/gcs.py @@ -25,7 +25,7 @@ from typing import Callable, List, Optional from airflow import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.sensors.base_sensor_operator import BaseSensorOperator from airflow.utils.decorators import apply_defaults @@ -66,7 +66,7 @@ def __init__(self, def poke(self, context): self.log.info('Sensor checks existence of : %s, %s', self.bucket, self.object) - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.google_cloud_conn_id, delegate_to=self.delegate_to) return hook.exists(self.bucket, self.object) @@ -123,7 +123,7 @@ def __init__(self, def poke(self, context): self.log.info('Sensor checks existence of : %s, %s', self.bucket, self.object) - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.google_cloud_conn_id, delegate_to=self.delegate_to) return hook.is_updated_after(self.bucket, self.object, self.ts_func(context)) @@ -170,7 +170,7 @@ def __init__(self, def poke(self, context): self.log.info('Sensor checks existence of objects: %s, %s', self.bucket, self.prefix) - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.google_cloud_conn_id, delegate_to=self.delegate_to) self._matches = hook.list(self.bucket, prefix=self.prefix) @@ -315,5 +315,5 @@ def is_bucket_updated(self, current_num_objects: int) -> bool: return False def poke(self, context): - hook = GoogleCloudStorageHook() + hook = GcsHook() return self.is_bucket_updated(len(hook.list(self.bucket, prefix=self.prefix))) diff --git a/airflow/gcp/utils/mlengine_operator_utils.py b/airflow/gcp/utils/mlengine_operator_utils.py index f838879ad6de5..235e8af4527bd 100644 --- a/airflow/gcp/utils/mlengine_operator_utils.py +++ b/airflow/gcp/utils/mlengine_operator_utils.py @@ -29,7 +29,7 @@ import dill from airflow.exceptions import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.gcp.operators.dataflow import DataFlowPythonOperator from airflow.gcp.operators.mlengine import MLEngineBatchPredictionOperator from airflow.operators.python_operator import PythonOperator @@ -244,7 +244,7 @@ def apply_validate_fn(*args, **kwargs): raise ValueError("Wrong format prediction_path: {}".format(prediction_path)) summary = os.path.join(obj.strip("/"), "prediction.summary.json") - gcs_hook = GoogleCloudStorageHook() + gcs_hook = GcsHook() summary = json.loads(gcs_hook.download(bucket, summary)) return validate_fn(summary) diff --git a/airflow/operators/adls_to_gcs.py b/airflow/operators/adls_to_gcs.py index b24d14305c1e5..83aa18a343325 100644 --- a/airflow/operators/adls_to_gcs.py +++ b/airflow/operators/adls_to_gcs.py @@ -27,7 +27,7 @@ from airflow.contrib.hooks.azure_data_lake_hook import AzureDataLakeHook from airflow.contrib.operators.adls_list_operator import AzureDataLakeStorageListOperator -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook, _parse_gcs_url +from airflow.gcp.hooks.gcs import GcsHook, _parse_gcs_url from airflow.utils.decorators import apply_defaults @@ -134,7 +134,7 @@ def __init__(self, def execute(self, context): # use the super to list all files in an Azure Data Lake path files = super().execute(context) - g_hook = GoogleCloudStorageHook( + g_hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) diff --git a/airflow/operators/cassandra_to_gcs.py b/airflow/operators/cassandra_to_gcs.py index 2b0c3943652f0..91f0ae3ea80a3 100644 --- a/airflow/operators/cassandra_to_gcs.py +++ b/airflow/operators/cassandra_to_gcs.py @@ -33,7 +33,7 @@ from airflow.contrib.hooks.cassandra_hook import CassandraHook from airflow.exceptions import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -217,7 +217,7 @@ def _write_local_schema_file(self, cursor): return {self.schema_filename: tmp_schema_file_handle} def _upload_to_gcs(self, files_to_upload): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) for object, tmp_file_handle in files_to_upload.items(): diff --git a/airflow/operators/gcs_to_bq.py b/airflow/operators/gcs_to_bq.py index 522ab7e7d67b4..ca182c2a9c3f6 100644 --- a/airflow/operators/gcs_to_bq.py +++ b/airflow/operators/gcs_to_bq.py @@ -24,7 +24,7 @@ from airflow import AirflowException from airflow.gcp.hooks.bigquery import BigQueryHook -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -238,7 +238,7 @@ def execute(self, context): if not self.schema_fields: if self.schema_object and self.source_format != 'DATASTORE_BACKUP': - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( google_cloud_storage_conn_id=self.google_cloud_storage_conn_id, delegate_to=self.delegate_to) schema_fields = json.loads(gcs_hook.download( diff --git a/airflow/operators/gcs_to_gcs.py b/airflow/operators/gcs_to_gcs.py index 944892a62616e..2f28dee216039 100644 --- a/airflow/operators/gcs_to_gcs.py +++ b/airflow/operators/gcs_to_gcs.py @@ -23,7 +23,7 @@ from typing import Optional from airflow.exceptions import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -156,7 +156,7 @@ def __init__(self, def execute(self, context): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) @@ -284,7 +284,7 @@ def __init__( self.delegate_to = delegate_to def execute(self, context): - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) diff --git a/airflow/operators/gcs_to_s3.py b/airflow/operators/gcs_to_s3.py index b4b41c5fac816..d782b5b4c9726 100644 --- a/airflow/operators/gcs_to_s3.py +++ b/airflow/operators/gcs_to_s3.py @@ -21,7 +21,7 @@ """ import warnings -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.gcp.operators.gcs import GoogleCloudStorageListOperator from airflow.hooks.S3_hook import S3Hook from airflow.utils.decorators import apply_defaults @@ -133,7 +133,7 @@ def execute(self, context): files = list(set(files) - set(existing_files)) if files: - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) diff --git a/airflow/operators/gcs_to_sftp.py b/airflow/operators/gcs_to_sftp.py index 90a5a7a6c9179..0eee7207661a1 100644 --- a/airflow/operators/gcs_to_sftp.py +++ b/airflow/operators/gcs_to_sftp.py @@ -25,7 +25,7 @@ from airflow import AirflowException from airflow.contrib.hooks.sftp_hook import SFTPHook -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -93,7 +93,7 @@ def __init__( self.sftp_dirs = None def execute(self, context): - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) @@ -132,7 +132,7 @@ def execute(self, context): def _copy_single_object( self, - gcs_hook: GoogleCloudStorageHook, + gcs_hook: GcsHook, sftp_hook: SFTPHook, source_object: str, destination_path: str, diff --git a/airflow/operators/local_to_gcs.py b/airflow/operators/local_to_gcs.py index a9de78896fef7..5657ae6ba174d 100644 --- a/airflow/operators/local_to_gcs.py +++ b/airflow/operators/local_to_gcs.py @@ -21,7 +21,7 @@ """ import warnings -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -84,7 +84,7 @@ def execute(self, context): """ Uploads the file to Google cloud storage """ - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) diff --git a/airflow/operators/sql_to_gcs.py b/airflow/operators/sql_to_gcs.py index a232c3b387962..b7b08a6c61aad 100644 --- a/airflow/operators/sql_to_gcs.py +++ b/airflow/operators/sql_to_gcs.py @@ -27,7 +27,7 @@ import unicodecsv as csv -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -268,7 +268,7 @@ def _upload_to_gcs(self, files_to_upload): Upload all of the file splits (and optionally the schema .json file) to Google cloud storage. """ - hook = GoogleCloudStorageHook( + hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) for tmp_file in files_to_upload: diff --git a/airflow/providers/google/marketing_platform/operators/campaign_manager.py b/airflow/providers/google/marketing_platform/operators/campaign_manager.py index e27b1c54bdcd2..9b99b38315967 100644 --- a/airflow/providers/google/marketing_platform/operators/campaign_manager.py +++ b/airflow/providers/google/marketing_platform/operators/campaign_manager.py @@ -26,7 +26,7 @@ from googleapiclient import http from airflow import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models.baseoperator import BaseOperator from airflow.providers.google.marketing_platform.hooks.campaign_manager import GoogleCampaignManagerHook from airflow.utils.decorators import apply_defaults @@ -210,7 +210,7 @@ def execute(self, context: Dict): delegate_to=self.delegate_to, api_version=self.api_version, ) - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) # Get name of the report diff --git a/airflow/providers/google/marketing_platform/operators/display_video.py b/airflow/providers/google/marketing_platform/operators/display_video.py index 3bc3dbb0754ce..5536b04b75c43 100644 --- a/airflow/providers/google/marketing_platform/operators/display_video.py +++ b/airflow/providers/google/marketing_platform/operators/display_video.py @@ -26,7 +26,7 @@ from urllib.parse import urlparse from airflow import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models.baseoperator import BaseOperator from airflow.providers.google.marketing_platform.hooks.display_video import GoogleDisplayVideo360Hook from airflow.utils.decorators import apply_defaults @@ -241,7 +241,7 @@ def execute(self, context: Dict): delegate_to=self.delegate_to, api_version=self.api_version, ) - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) diff --git a/airflow/providers/google/marketing_platform/operators/search_ads.py b/airflow/providers/google/marketing_platform/operators/search_ads.py index dcd15f4efeddb..8b655d057aadb 100644 --- a/airflow/providers/google/marketing_platform/operators/search_ads.py +++ b/airflow/providers/google/marketing_platform/operators/search_ads.py @@ -23,7 +23,7 @@ from typing import Any, Dict, Optional from airflow import AirflowException -from airflow.gcp.hooks.gcs import GoogleCloudStorageHook +from airflow.gcp.hooks.gcs import GcsHook from airflow.models.baseoperator import BaseOperator from airflow.providers.google.marketing_platform.hooks.search_ads import GoogleSearchAdsHook from airflow.utils.decorators import apply_defaults @@ -170,7 +170,7 @@ def execute(self, context: Dict): api_version=self.api_version, ) - gcs_hook = GoogleCloudStorageHook( + gcs_hook = GcsHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) diff --git a/airflow/utils/log/gcs_task_handler.py b/airflow/utils/log/gcs_task_handler.py index e1e987382834b..b00404afec239 100644 --- a/airflow/utils/log/gcs_task_handler.py +++ b/airflow/utils/log/gcs_task_handler.py @@ -46,13 +46,13 @@ def __init__(self, base_log_folder, gcs_log_folder, filename_template): def hook(self): remote_conn_id = conf.get('core', 'REMOTE_LOG_CONN_ID') try: - from airflow.gcp.hooks.gcs import GoogleCloudStorageHook - return GoogleCloudStorageHook( + from airflow.gcp.hooks.gcs import GcsHook + return GcsHook( google_cloud_storage_conn_id=remote_conn_id ) except Exception as e: self.log.error( - 'Could not create a GoogleCloudStorageHook with connection id ' + 'Could not create a GcsHook with connection id ' '"%s". %s\n\nPlease make sure that airflow[gcp] is installed ' 'and the GCS connection exists.', remote_conn_id, str(e) ) diff --git a/tests/contrib/operators/test_gcs_to_gdrive.py b/tests/contrib/operators/test_gcs_to_gdrive.py index 0f4d4cb53cd83..16b8f8db5b8fb 100644 --- a/tests/contrib/operators/test_gcs_to_gdrive.py +++ b/tests/contrib/operators/test_gcs_to_gdrive.py @@ -26,7 +26,7 @@ class TestGcsToGDriveOperator(unittest.TestCase): - @mock.patch(MODULE + ".GoogleCloudStorageHook") + @mock.patch(MODULE + ".GcsHook") @mock.patch(MODULE + ".GoogleDriveHook") @mock.patch(MODULE + ".tempfile.NamedTemporaryFile") def test_should_copy_single_file(self, mock_named_temporary_file, mock_gdrive, mock_gcs_hook): @@ -61,7 +61,7 @@ def test_should_copy_single_file(self, mock_named_temporary_file, mock_gdrive, m ) # - @mock.patch(MODULE + ".GoogleCloudStorageHook") + @mock.patch(MODULE + ".GcsHook") @mock.patch(MODULE + ".GoogleDriveHook") @mock.patch(MODULE + ".tempfile.NamedTemporaryFile") def test_should_copy_files(self, mock_named_temporary_file, mock_gdrive, mock_gcs_hook): @@ -97,7 +97,7 @@ def test_should_copy_files(self, mock_named_temporary_file, mock_gdrive, mock_gc ] ) - @mock.patch(MODULE + ".GoogleCloudStorageHook") + @mock.patch(MODULE + ".GcsHook") @mock.patch(MODULE + ".GoogleDriveHook") @mock.patch(MODULE + ".tempfile.NamedTemporaryFile") def test_should_move_files(self, mock_named_temporary_file, mock_gdrive, mock_gcs_hook): @@ -135,7 +135,7 @@ def test_should_move_files(self, mock_named_temporary_file, mock_gdrive, mock_gc ] ) - @mock.patch(MODULE + ".GoogleCloudStorageHook") + @mock.patch(MODULE + ".GcsHook") @mock.patch(MODULE + ".GoogleDriveHook") @mock.patch(MODULE + ".tempfile.NamedTemporaryFile") def test_should_raise_exception_on_multiple_wildcard( diff --git a/tests/contrib/operators/test_s3_to_gcs_operator.py b/tests/contrib/operators/test_s3_to_gcs_operator.py index 42821a755fe88..8e904d994a867 100644 --- a/tests/contrib/operators/test_s3_to_gcs_operator.py +++ b/tests/contrib/operators/test_s3_to_gcs_operator.py @@ -54,7 +54,7 @@ def test_init(self): @mock.patch('airflow.contrib.operators.s3_to_gcs_operator.S3Hook') @mock.patch('airflow.contrib.operators.s3_list_operator.S3Hook') @mock.patch( - 'airflow.contrib.operators.s3_to_gcs_operator.GoogleCloudStorageHook') + 'airflow.contrib.operators.s3_to_gcs_operator.GcsHook') def test_execute(self, gcs_mock_hook, s3_one_mock_hook, s3_two_mock_hook): """Test the execute function when the run is successful.""" @@ -89,7 +89,7 @@ def test_execute(self, gcs_mock_hook, s3_one_mock_hook, s3_two_mock_hook): @mock.patch('airflow.contrib.operators.s3_to_gcs_operator.S3Hook') @mock.patch('airflow.contrib.operators.s3_list_operator.S3Hook') @mock.patch( - 'airflow.contrib.operators.s3_to_gcs_operator.GoogleCloudStorageHook') + 'airflow.contrib.operators.s3_to_gcs_operator.GcsHook') def test_execute_with_gzip(self, gcs_mock_hook, s3_one_mock_hook, s3_two_mock_hook): """Test the execute function when the run is successful.""" diff --git a/tests/gcp/hooks/test_gcs.py b/tests/gcp/hooks/test_gcs.py index bdb7cacd18a34..993aed5173795 100644 --- a/tests/gcp/hooks/test_gcs.py +++ b/tests/gcp/hooks/test_gcs.py @@ -71,7 +71,7 @@ def setUp(self): GCS_STRING.format('GoogleCloudBaseHook.__init__'), new=mock_base_gcp_hook_default_project_id, ): - self.gcs_hook = gcs.GoogleCloudStorageHook( + self.gcs_hook = gcs.GcsHook( google_cloud_storage_conn_id='test') @mock.patch( @@ -90,7 +90,7 @@ def test_storage_client_creation(self, mock_get_connetion, mock_get_creds_and_project_id, mock_client_info): - hook = gcs.GoogleCloudStorageHook() + hook = gcs.GcsHook() result = hook.get_conn() # test that Storage Client is called with required arguments mock_client.assert_called_once_with( @@ -99,7 +99,7 @@ def test_storage_client_creation(self, project="PROJECT_ID") self.assertEqual(mock_client.return_value, result) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_exists(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -119,7 +119,7 @@ def test_exists(self, mock_service): blob_object.assert_called_once_with(blob_name=test_object) exists_method.assert_called_once_with() - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_exists_nonexisting_object(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -136,7 +136,7 @@ def test_exists_nonexisting_object(self, mock_service): # Then self.assertFalse(response) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_is_updated_after(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -155,7 +155,7 @@ def test_is_updated_after(self, mock_service): self.assertTrue(response) @mock.patch('google.cloud.storage.Bucket') - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_copy(self, mock_service, mock_bucket): source_bucket = 'test-source-bucket' source_object = 'test-source-object' @@ -244,7 +244,7 @@ def test_copy_empty_source_object(self): ) @mock.patch('google.cloud.storage.Bucket') - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_rewrite(self, mock_service, mock_bucket): source_bucket = 'test-source-bucket' source_object = 'test-source-object' @@ -307,7 +307,7 @@ def test_rewrite_empty_source_object(self): ) @mock.patch('google.cloud.storage.Bucket') - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_delete(self, mock_service, mock_bucket): test_bucket = 'test_bucket' test_object = 'test_object' @@ -323,7 +323,7 @@ def test_delete(self, mock_service, mock_bucket): object_name=test_object) self.assertIsNone(response) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_delete_nonexisting_object(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -336,7 +336,7 @@ def test_delete_nonexisting_object(self, mock_service): with self.assertRaises(exceptions.NotFound): self.gcs_hook.delete(bucket_name=test_bucket, object_name=test_object) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_object_get_size(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -351,7 +351,7 @@ def test_object_get_size(self, mock_service): self.assertEqual(response, returned_file_size) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_object_get_crc32c(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -366,7 +366,7 @@ def test_object_get_crc32c(self, mock_service): self.assertEqual(response, returned_file_crc32c) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_object_get_md5hash(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -382,7 +382,7 @@ def test_object_get_md5hash(self, mock_service): self.assertEqual(response, returned_file_md5hash) @mock.patch('google.cloud.storage.Bucket') - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_create_bucket(self, mock_service, mock_bucket): test_bucket = 'test_bucket' test_project = 'test-project' @@ -417,7 +417,7 @@ def test_create_bucket(self, mock_service, mock_bucket): ) @mock.patch('google.cloud.storage.Bucket') - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_create_bucket_with_resource(self, mock_service, mock_bucket): test_bucket = 'test_bucket' test_project = 'test-project' @@ -454,7 +454,7 @@ def test_create_bucket_with_resource(self, mock_service, mock_bucket): ) @mock.patch('google.cloud.storage.Bucket.blob') - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_compose(self, mock_service, mock_blob): test_bucket = 'test_bucket' test_source_objects = ['test_object_1', 'test_object_2', 'test_object_3'] @@ -476,7 +476,7 @@ def test_compose(self, mock_service, mock_blob): mock_blob(blob_name=source_object) for source_object in test_source_objects ]) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_compose_with_empty_source_objects(self, mock_service): # pylint:disable=unused-argument test_bucket = 'test_bucket' test_source_objects = [] @@ -494,7 +494,7 @@ def test_compose_with_empty_source_objects(self, mock_service): # pylint:disabl 'source_objects cannot be empty.' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_compose_without_bucket(self, mock_service): # pylint:disable=unused-argument test_bucket = None test_source_objects = ['test_object_1', 'test_object_2', 'test_object_3'] @@ -512,7 +512,7 @@ def test_compose_without_bucket(self, mock_service): # pylint:disable=unused-ar 'bucket_name and destination_object cannot be empty.' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_compose_without_destination_object(self, mock_service): # pylint:disable=unused-argument test_bucket = 'test_bucket' test_source_objects = ['test_object_1', 'test_object_2', 'test_object_3'] @@ -530,7 +530,7 @@ def test_compose_without_destination_object(self, mock_service): # pylint:disab 'bucket_name and destination_object cannot be empty.' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_download_as_string(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -547,7 +547,7 @@ def test_download_as_string(self, mock_service): self.assertEqual(response, test_object_bytes) download_method.assert_called_once_with() - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_download_to_file(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -572,7 +572,7 @@ def test_download_to_file(self, mock_service): class TestGoogleCloudStorageHookUpload(unittest.TestCase): def setUp(self): with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__')): - self.gcs_hook = gcs.GoogleCloudStorageHook( + self.gcs_hook = gcs.GcsHook( google_cloud_storage_conn_id='test' ) @@ -586,7 +586,7 @@ def setUp(self): def tearDown(self): os.unlink(self.testfile.name) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_upload_file(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -603,7 +603,7 @@ def test_upload_file(self, mock_service): content_type='application/octet-stream' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_upload_file_gzip(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -614,7 +614,7 @@ def test_upload_file_gzip(self, mock_service): gzip=True) self.assertFalse(os.path.exists(self.testfile.name + '.gz')) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_upload_data_str(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -631,7 +631,7 @@ def test_upload_data_str(self, mock_service): content_type='text/plain' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_upload_data_bytes(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -648,7 +648,7 @@ def test_upload_data_bytes(self, mock_service): content_type='text/plain' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_upload_data_str_gzip(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -673,7 +673,7 @@ def test_upload_data_str_gzip(self, mock_service): content_type='text/plain' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_upload_data_bytes_gzip(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -697,7 +697,7 @@ def test_upload_data_bytes_gzip(self, mock_service): content_type='text/plain' ) - @mock.patch(GCS_STRING.format('GoogleCloudStorageHook.get_conn')) + @mock.patch(GCS_STRING.format('GcsHook.get_conn')) def test_upload_exceptions(self, mock_service): test_bucket = 'test_bucket' test_object = 'test_object' @@ -722,12 +722,12 @@ def setUp(self): with mock.patch( GCS_STRING.format("GoogleCloudBaseHook.__init__"), new=mock_base_gcp_hook_default_project_id ): - self.gcs_hook = gcs.GoogleCloudStorageHook(google_cloud_storage_conn_id="test") + self.gcs_hook = gcs.GcsHook(google_cloud_storage_conn_id="test") - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_do_nothing_when_buckets_is_empty( self, mock_get_conn, mock_delete, mock_rewrite, mock_copy ): @@ -747,10 +747,10 @@ def test_should_do_nothing_when_buckets_is_empty( mock_rewrite.assert_not_called() mock_copy.assert_not_called() - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_append_slash_to_object_if_missing( self, mock_get_conn, mock_delete, mock_rewrite, mock_copy ): @@ -769,10 +769,10 @@ def test_should_append_slash_to_object_if_missing( source_bucket.list_blobs.assert_called_once_with(delimiter=None, prefix="SOURCE_OBJECT/") destination_bucket.list_blobs.assert_called_once_with(delimiter=None, prefix="DESTINATION_OBJECT/") - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_copy_files(self, mock_get_conn, mock_delete, mock_rewrite, mock_copy): # mock_get_conn.return_value = source_bucket = self._create_bucket(name="SOURCE_BUCKET") @@ -804,10 +804,10 @@ def test_should_copy_files(self, mock_get_conn, mock_delete, mock_rewrite, mock_ any_order=True, ) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_copy_files_non_recursive(self, mock_get_conn, mock_delete, mock_rewrite, mock_copy): # mock_get_conn.return_value = source_bucket = self._create_bucket(name="SOURCE_BUCKET") @@ -822,10 +822,10 @@ def test_should_copy_files_non_recursive(self, mock_get_conn, mock_delete, mock_ source_bucket.list_blobs.assert_called_once_with(delimiter='/', prefix=None) destination_bucket.list_blobs.assert_called_once_with(delimiter='/', prefix=None) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_copy_files_to_subdirectory(self, mock_get_conn, mock_delete, mock_rewrite, mock_copy): # mock_get_conn.return_value = source_bucket = self._create_bucket(name="SOURCE_BUCKET") @@ -859,10 +859,10 @@ def test_should_copy_files_to_subdirectory(self, mock_get_conn, mock_delete, moc any_order=True, ) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_copy_files_from_subdirectory(self, mock_get_conn, mock_delete, mock_rewrite, mock_copy): # mock_get_conn.return_value = source_bucket = self._create_bucket(name="SOURCE_BUCKET") @@ -896,10 +896,10 @@ def test_should_copy_files_from_subdirectory(self, mock_get_conn, mock_delete, m any_order=True, ) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_overwrite_files(self, mock_get_conn, mock_delete, mock_rewrite, mock_copy): # mock_get_conn.return_value = source_bucket = self._create_bucket(name="SOURCE_BUCKET") @@ -936,10 +936,10 @@ def test_should_overwrite_files(self, mock_get_conn, mock_delete, mock_rewrite, ) mock_copy.assert_not_called() - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_overwrite_files_to_subdirectory( self, mock_get_conn, mock_delete, mock_rewrite, mock_copy ): @@ -981,10 +981,10 @@ def test_should_overwrite_files_to_subdirectory( ) mock_copy.assert_not_called() - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_overwrite_files_from_subdirectory( self, mock_get_conn, mock_delete, mock_rewrite, mock_copy ): @@ -1026,10 +1026,10 @@ def test_should_overwrite_files_from_subdirectory( ) mock_copy.assert_not_called() - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_delete_extra_files(self, mock_get_conn, mock_delete, mock_rewrite, mock_copy): # mock_get_conn.return_value = source_bucket = self._create_bucket(name="SOURCE_BUCKET") @@ -1050,10 +1050,10 @@ def test_should_delete_extra_files(self, mock_get_conn, mock_delete, mock_rewrit mock_rewrite.assert_not_called() mock_copy.assert_not_called() - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_not_delete_extra_files_when_delete_extra_files_is_disabled( self, mock_get_conn, mock_delete, mock_rewrite, mock_copy ): @@ -1073,10 +1073,10 @@ def test_should_not_delete_extra_files_when_delete_extra_files_is_disabled( mock_rewrite.assert_not_called() mock_copy.assert_not_called() - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.copy")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.rewrite")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.delete")) - @mock.patch(GCS_STRING.format("GoogleCloudStorageHook.get_conn")) + @mock.patch(GCS_STRING.format("GcsHook.copy")) + @mock.patch(GCS_STRING.format("GcsHook.rewrite")) + @mock.patch(GCS_STRING.format("GcsHook.delete")) + @mock.patch(GCS_STRING.format("GcsHook.get_conn")) def test_should_not_overwrite_when_overwrite_is_disabled( self, mock_get_conn, mock_delete, mock_rewrite, mock_copy ): diff --git a/tests/gcp/operators/test_gcs.py b/tests/gcp/operators/test_gcs.py index 68b1cf93dbd8c..ca507cd2bcdfe 100644 --- a/tests/gcp/operators/test_gcs.py +++ b/tests/gcp/operators/test_gcs.py @@ -39,7 +39,7 @@ class TestGoogleCloudStorageCreateBucket(unittest.TestCase): - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_execute(self, mock_hook): operator = GoogleCloudStorageCreateBucketOperator( task_id=TASK_ID, @@ -61,7 +61,7 @@ def test_execute(self, mock_hook): class TestGoogleCloudStorageAcl(unittest.TestCase): - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_bucket_create_acl(self, mock_hook): operator = GoogleCloudStorageBucketCreateAclEntryOperator( bucket="test-bucket", @@ -78,7 +78,7 @@ def test_bucket_create_acl(self, mock_hook): user_project="test-user-project" ) - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_object_create_acl(self, mock_hook): operator = GoogleCloudStorageObjectCreateAclEntryOperator( bucket="test-bucket", @@ -101,7 +101,7 @@ def test_object_create_acl(self, mock_hook): class TestGoogleCloudStorageDeleteOperator(unittest.TestCase): - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_delete_objects(self, mock_hook): operator = GoogleCloudStorageDeleteOperator(task_id=TASK_ID, bucket_name=TEST_BUCKET, @@ -117,7 +117,7 @@ def test_delete_objects(self, mock_hook): any_order=True ) - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_delete_prefix(self, mock_hook): mock_hook.return_value.list.return_value = MOCK_FILES[1:3] operator = GoogleCloudStorageDeleteOperator(task_id=TASK_ID, @@ -139,7 +139,7 @@ def test_delete_prefix(self, mock_hook): class TestGoogleCloudStorageDownloadOperator(unittest.TestCase): - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_execute(self, mock_hook): operator = GoogleCloudStorageDownloadOperator(task_id=TASK_ID, bucket=TEST_BUCKET, @@ -154,7 +154,7 @@ def test_execute(self, mock_hook): class TestGoogleCloudStorageListOperator(unittest.TestCase): - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_execute(self, mock_hook): mock_hook.return_value.list.return_value = MOCK_FILES @@ -173,7 +173,7 @@ def test_execute(self, mock_hook): class TestGcsFileTransformOperator(unittest.TestCase): @mock.patch('airflow.gcp.operators.gcs.NamedTemporaryFile') @mock.patch('airflow.gcp.operators.gcs.subprocess') - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') def test_execute(self, mock_hook, mock_subprocess, mock_tempfile): source_bucket = TEST_BUCKET source_object = "test.txt" diff --git a/tests/gcp/operators/test_mlengine_utils.py b/tests/gcp/operators/test_mlengine_utils.py index b58e78be28722..1d6f3ebf17351 100644 --- a/tests/gcp/operators/test_mlengine_utils.py +++ b/tests/gcp/operators/test_mlengine_utils.py @@ -114,7 +114,7 @@ def test_successful_run(self): ['-m'], py_interpreter='python3') with patch('airflow.gcp.utils.mlengine_operator_utils.' - 'GoogleCloudStorageHook') as mock_gcs_hook: + 'GcsHook') as mock_gcs_hook: hook_instance = mock_gcs_hook.return_value hook_instance.download.return_value = '{"err": 0.9, "count": 9}' result = validate.execute({}) diff --git a/tests/gcp/operators/test_text_to_speech.py b/tests/gcp/operators/test_text_to_speech.py index db115fbfdcc18..790a80e6d0905 100644 --- a/tests/gcp/operators/test_text_to_speech.py +++ b/tests/gcp/operators/test_text_to_speech.py @@ -35,7 +35,7 @@ class TestGcpTextToSpeech(unittest.TestCase): - @patch("airflow.gcp.operators.text_to_speech.GoogleCloudStorageHook") + @patch("airflow.gcp.operators.text_to_speech.GcsHook") @patch("airflow.gcp.operators.text_to_speech.GCPTextToSpeechHook") def test_synthesize_text_green_path(self, mock_text_to_speech_hook, mock_gcp_hook): mocked_response = Mock() @@ -73,7 +73,7 @@ def test_synthesize_text_green_path(self, mock_text_to_speech_hook, mock_gcp_hoo ("target_filename", INPUT, VOICE, AUDIO_CONFIG, TARGET_BUCKET_NAME, ""), ] ) - @patch("airflow.gcp.operators.text_to_speech.GoogleCloudStorageHook") + @patch("airflow.gcp.operators.text_to_speech.GcsHook") @patch("airflow.gcp.operators.text_to_speech.GCPTextToSpeechHook") def test_missing_arguments( self, diff --git a/tests/gcp/sensors/test_gcs.py b/tests/gcp/sensors/test_gcs.py index 36b1469f1bad6..a9a4314dfeaa8 100644 --- a/tests/gcp/sensors/test_gcs.py +++ b/tests/gcp/sensors/test_gcs.py @@ -61,7 +61,7 @@ def reset(dag_id=TEST_DAG_ID): class TestGoogleCloudStorageObjectSensor(TestCase): - @mock.patch("airflow.gcp.sensors.gcs.GoogleCloudStorageHook") + @mock.patch("airflow.gcp.sensors.gcs.GcsHook") def test_should_pass_argument_to_hook(self, mock_hook): task = GoogleCloudStorageObjectSensor( task_id="task-id", @@ -107,7 +107,7 @@ def test_should_support_cron(self): class TestGoogleCloudStorageObjectUpdatedSensor(TestCase): - @mock.patch("airflow.gcp.sensors.gcs.GoogleCloudStorageHook") + @mock.patch("airflow.gcp.sensors.gcs.GcsHook") def test_should_pass_argument_to_hook(self, mock_hook): task = GoogleCloudStorageObjectUpdatedSensor( task_id="task-id", @@ -128,7 +128,7 @@ def test_should_pass_argument_to_hook(self, mock_hook): class TestGoogleCloudStoragePrefixSensor(TestCase): - @mock.patch("airflow.gcp.sensors.gcs.GoogleCloudStorageHook") + @mock.patch("airflow.gcp.sensors.gcs.GcsHook") def test_should_pass_arguments_to_hook(self, mock_hook): task = GoogleCloudStoragePrefixSensor( task_id="task-id", @@ -147,7 +147,7 @@ def test_should_pass_arguments_to_hook(self, mock_hook): mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET, prefix=TEST_PREFIX) self.assertEqual(True, result) - @mock.patch("airflow.gcp.sensors.gcs.GoogleCloudStorageHook") + @mock.patch("airflow.gcp.sensors.gcs.GcsHook") def test_should_return_false_on_empty_list(self, mock_hook): task = GoogleCloudStoragePrefixSensor( task_id="task-id", @@ -161,7 +161,7 @@ def test_should_return_false_on_empty_list(self, mock_hook): self.assertEqual(False, result) - @mock.patch('airflow.gcp.sensors.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.sensors.gcs.GcsHook') def test_execute(self, mock_hook): task = GoogleCloudStoragePrefixSensor( task_id="task-id", @@ -182,7 +182,7 @@ def test_execute(self, mock_hook): mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET, prefix=TEST_PREFIX) self.assertEqual(response, generated_messages) - @mock.patch('airflow.gcp.sensors.gcs.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.sensors.gcs.GcsHook') def test_execute_timeout(self, mock_hook): task = GoogleCloudStoragePrefixSensor( task_id="task-id", diff --git a/tests/operators/test_adls_to_gcs_operator.py b/tests/operators/test_adls_to_gcs_operator.py index df0ac2c667060..5fbab1fea499f 100644 --- a/tests/operators/test_adls_to_gcs_operator.py +++ b/tests/operators/test_adls_to_gcs_operator.py @@ -54,7 +54,7 @@ def test_init(self): @mock.patch('airflow.operators.adls_to_gcs.AzureDataLakeHook') @mock.patch('airflow.contrib.operators.adls_list_operator.AzureDataLakeHook') @mock.patch( - 'airflow.operators.adls_to_gcs.GoogleCloudStorageHook') + 'airflow.operators.adls_to_gcs.GcsHook') def test_execute(self, gcs_mock_hook, adls_one_mock_hook, adls_two_mock_hook): """Test the execute function when the run is successful.""" @@ -103,7 +103,7 @@ def test_execute(self, gcs_mock_hook, adls_one_mock_hook, adls_two_mock_hook): @mock.patch('airflow.operators.adls_to_gcs.AzureDataLakeHook') @mock.patch('airflow.contrib.operators.adls_list_operator.AzureDataLakeHook') @mock.patch( - 'airflow.operators.adls_to_gcs.GoogleCloudStorageHook') + 'airflow.operators.adls_to_gcs.GcsHook') def test_execute_with_gzip(self, gcs_mock_hook, adls_one_mock_hook, adls_two_mock_hook): """Test the execute function when the run is successful.""" diff --git a/tests/operators/test_cassandra_to_gcs.py b/tests/operators/test_cassandra_to_gcs.py index 1a10228f1035f..7c2edb9ca7321 100644 --- a/tests/operators/test_cassandra_to_gcs.py +++ b/tests/operators/test_cassandra_to_gcs.py @@ -30,7 +30,7 @@ class TestCassandraToGCS(unittest.TestCase): @mock.patch("airflow.operators.cassandra_to_gcs.NamedTemporaryFile") @mock.patch( - "airflow.operators.cassandra_to_gcs.GoogleCloudStorageHook.upload" + "airflow.operators.cassandra_to_gcs.GcsHook.upload" ) @mock.patch("airflow.operators.cassandra_to_gcs.CassandraHook") def test_execute(self, mock_hook, mock_upload, mock_tempfile): diff --git a/tests/operators/test_gcs_to_gcs.py b/tests/operators/test_gcs_to_gcs.py index ee7d1c07f3488..ad58ab0ef345c 100644 --- a/tests/operators/test_gcs_to_gcs.py +++ b/tests/operators/test_gcs_to_gcs.py @@ -56,7 +56,7 @@ class TestGoogleCloudStorageToCloudStorageOperator(unittest.TestCase): Also tests the destionation_object as prefix when the wildcard is used. """ - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_no_prefix(self, mock_hook): operator = GoogleCloudStorageToGoogleCloudStorageOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, @@ -68,7 +68,7 @@ def test_execute_no_prefix(self, mock_hook): TEST_BUCKET, prefix="", delimiter="test_object" ) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_no_suffix(self, mock_hook): operator = GoogleCloudStorageToGoogleCloudStorageOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, @@ -80,7 +80,7 @@ def test_execute_no_suffix(self, mock_hook): TEST_BUCKET, prefix="test_object", delimiter="" ) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_prefix_and_suffix(self, mock_hook): operator = GoogleCloudStorageToGoogleCloudStorageOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, @@ -94,7 +94,7 @@ def test_execute_prefix_and_suffix(self, mock_hook): # copy with wildcard - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_wildcard_with_destination_object(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -112,7 +112,7 @@ def test_execute_wildcard_with_destination_object(self, mock_hook): ] mock_hook.return_value.rewrite.assert_has_calls(mock_calls) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_wildcard_with_destination_object_retained_prefix(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -132,7 +132,7 @@ def test_execute_wildcard_with_destination_object_retained_prefix(self, mock_hoo ] mock_hook.return_value.rewrite.assert_has_calls(mock_calls_retained) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_wildcard_without_destination_object(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -149,7 +149,7 @@ def test_execute_wildcard_without_destination_object(self, mock_hook): ] mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_wildcard_empty_destination_object(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -167,7 +167,7 @@ def test_execute_wildcard_empty_destination_object(self, mock_hook): ] mock_hook.return_value.rewrite.assert_has_calls(mock_calls_empty) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_last_modified_time(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -185,7 +185,7 @@ def test_execute_last_modified_time(self, mock_hook): ] mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_wc_with_last_modified_time_with_all_true_cond(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST mock_hook.return_value.is_updated_after.side_effect = [True, True, True] @@ -204,7 +204,7 @@ def test_wc_with_last_modified_time_with_all_true_cond(self, mock_hook): ] mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_wc_with_last_modified_time_with_one_true_cond(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST mock_hook.return_value.is_updated_after.side_effect = [True, False, False] @@ -219,7 +219,7 @@ def test_wc_with_last_modified_time_with_one_true_cond(self, mock_hook): TEST_BUCKET, 'test_object/file1.txt', DESTINATION_BUCKET, 'test_object/file1.txt') - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_wc_with_no_last_modified_time(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -237,7 +237,7 @@ def test_wc_with_no_last_modified_time(self, mock_hook): ] mock_hook.return_value.rewrite.assert_has_calls(mock_calls_none) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_no_prefix_with_last_modified_time_with_true_cond(self, mock_hook): mock_hook.return_value.is_updated_after.return_value = True operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -251,7 +251,7 @@ def test_no_prefix_with_last_modified_time_with_true_cond(self, mock_hook): mock_hook.return_value.rewrite.assert_called_once_with( TEST_BUCKET, 'test_object.txt', DESTINATION_BUCKET, 'test_object.txt') - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_no_prefix_with_no_last_modified_time(self, mock_hook): operator = GoogleCloudStorageToGoogleCloudStorageOperator( task_id=TASK_ID, source_bucket=TEST_BUCKET, @@ -264,7 +264,7 @@ def test_execute_no_prefix_with_no_last_modified_time(self, mock_hook): mock_hook.return_value.rewrite.assert_called_once_with( TEST_BUCKET, 'test_object.txt', DESTINATION_BUCKET, 'test_object.txt') - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_no_prefix_with_last_modified_time_with_false_cond(self, mock_hook): mock_hook.return_value.is_updated_after.return_value = False operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -277,7 +277,7 @@ def test_no_prefix_with_last_modified_time_with_false_cond(self, mock_hook): operator.execute(None) mock_hook.return_value.rewrite.assert_not_called() - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_more_than_1_wildcard(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -294,7 +294,7 @@ def test_execute_more_than_1_wildcard(self, mock_hook): with self.assertRaisesRegex(AirflowException, error_msg): operator.execute(None) - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute_with_empty_destination_bucket(self, mock_hook): mock_hook.return_value.list.return_value = SOURCE_FILES_LIST operator = GoogleCloudStorageToGoogleCloudStorageOperator( @@ -314,7 +314,7 @@ def test_execute_with_empty_destination_bucket(self, mock_hook): class TestGoogleCloudStorageSync(unittest.TestCase): - @mock.patch('airflow.operators.gcs_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.gcs_to_gcs.GcsHook') def test_execute(self, mock_hook): task = GoogleCloudStorageSynchronizeBuckets( task_id="task-id", diff --git a/tests/operators/test_gcs_to_s3.py b/tests/operators/test_gcs_to_s3.py index 4e9d2d2a8df5a..3ec83fca88f7d 100644 --- a/tests/operators/test_gcs_to_s3.py +++ b/tests/operators/test_gcs_to_s3.py @@ -40,8 +40,8 @@ class TestGoogleCloudStorageToS3Operator(unittest.TestCase): # Test1: incremental behaviour (just some files missing) @mock_s3 - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') - @mock.patch('airflow.operators.gcs_to_s3.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') + @mock.patch('airflow.operators.gcs_to_s3.GcsHook') def test_execute_incremental(self, mock_hook, mock_hook2): mock_hook.return_value.list.return_value = MOCK_FILES mock_hook.return_value.download.return_value = b"testing" @@ -70,8 +70,8 @@ def test_execute_incremental(self, mock_hook, mock_hook2): # Test2: All the files are already in origin and destination without replace @mock_s3 - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') - @mock.patch('airflow.operators.gcs_to_s3.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') + @mock.patch('airflow.operators.gcs_to_s3.GcsHook') def test_execute_without_replace(self, mock_hook, mock_hook2): mock_hook.return_value.list.return_value = MOCK_FILES mock_hook.return_value.download.return_value = b"testing" @@ -100,8 +100,8 @@ def test_execute_without_replace(self, mock_hook, mock_hook2): # Test3: There are no files in destination bucket @mock_s3 - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') - @mock.patch('airflow.operators.gcs_to_s3.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') + @mock.patch('airflow.operators.gcs_to_s3.GcsHook') def test_execute(self, mock_hook, mock_hook2): mock_hook.return_value.list.return_value = MOCK_FILES mock_hook.return_value.download.return_value = b"testing" @@ -129,8 +129,8 @@ def test_execute(self, mock_hook, mock_hook2): # Test4: Destination and Origin are in sync but replace all files in destination @mock_s3 - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') - @mock.patch('airflow.operators.gcs_to_s3.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') + @mock.patch('airflow.operators.gcs_to_s3.GcsHook') def test_execute_with_replace(self, mock_hook, mock_hook2): mock_hook.return_value.list.return_value = MOCK_FILES mock_hook.return_value.download.return_value = b"testing" @@ -159,8 +159,8 @@ def test_execute_with_replace(self, mock_hook, mock_hook2): # Test5: Incremental sync with replace @mock_s3 - @mock.patch('airflow.gcp.operators.gcs.GoogleCloudStorageHook') - @mock.patch('airflow.operators.gcs_to_s3.GoogleCloudStorageHook') + @mock.patch('airflow.gcp.operators.gcs.GcsHook') + @mock.patch('airflow.operators.gcs_to_s3.GcsHook') def test_execute_incremental_with_replace(self, mock_hook, mock_hook2): mock_hook.return_value.list.return_value = MOCK_FILES mock_hook.return_value.download.return_value = b"testing" diff --git a/tests/operators/test_gcs_to_sftp.py b/tests/operators/test_gcs_to_sftp.py index 526d20a563881..f86b400effd04 100644 --- a/tests/operators/test_gcs_to_sftp.py +++ b/tests/operators/test_gcs_to_sftp.py @@ -46,7 +46,7 @@ # pylint: disable=unused-argument class TestGoogleCloudStorageToSFTPOperator(unittest.TestCase): - @mock.patch("airflow.operators.gcs_to_sftp.GoogleCloudStorageHook") + @mock.patch("airflow.operators.gcs_to_sftp.GcsHook") @mock.patch("airflow.operators.gcs_to_sftp.SFTPHook") def test_execute_copy_single_file(self, sftp_hook, gcs_hook): task = GoogleCloudStorageToSFTPOperator( @@ -76,7 +76,7 @@ def test_execute_copy_single_file(self, sftp_hook, gcs_hook): gcs_hook.return_value.delete.assert_not_called() - @mock.patch("airflow.operators.gcs_to_sftp.GoogleCloudStorageHook") + @mock.patch("airflow.operators.gcs_to_sftp.GcsHook") @mock.patch("airflow.operators.gcs_to_sftp.SFTPHook") def test_execute_move_single_file(self, sftp_hook, gcs_hook): task = GoogleCloudStorageToSFTPOperator( @@ -108,7 +108,7 @@ def test_execute_move_single_file(self, sftp_hook, gcs_hook): TEST_BUCKET, SOURCE_OBJECT_NO_WILDCARD ) - @mock.patch("airflow.operators.gcs_to_sftp.GoogleCloudStorageHook") + @mock.patch("airflow.operators.gcs_to_sftp.GcsHook") @mock.patch("airflow.operators.gcs_to_sftp.SFTPHook") def test_execute_copy_with_wildcard(self, sftp_hook, gcs_hook): gcs_hook.return_value.list.return_value = SOURCE_FILES_LIST[:2] @@ -135,7 +135,7 @@ def test_execute_copy_with_wildcard(self, sftp_hook, gcs_hook): self.assertEqual(call_two[1]["bucket_name"], TEST_BUCKET) self.assertEqual(call_two[1]["object_name"], "test_object/file2.txt") - @mock.patch("airflow.operators.gcs_to_sftp.GoogleCloudStorageHook") + @mock.patch("airflow.operators.gcs_to_sftp.GcsHook") @mock.patch("airflow.operators.gcs_to_sftp.SFTPHook") def test_execute_move_with_wildcard(self, sftp_hook, gcs_hook): gcs_hook.return_value.list.return_value = SOURCE_FILES_LIST[:2] @@ -159,7 +159,7 @@ def test_execute_move_with_wildcard(self, sftp_hook, gcs_hook): self.assertEqual(call_one[0], (TEST_BUCKET, "test_object/file1.txt")) self.assertEqual(call_two[0], (TEST_BUCKET, "test_object/file2.txt")) - @mock.patch("airflow.operators.gcs_to_sftp.GoogleCloudStorageHook") + @mock.patch("airflow.operators.gcs_to_sftp.GcsHook") @mock.patch("airflow.operators.gcs_to_sftp.SFTPHook") def test_execute_more_than_one_wildcard_exception(self, sftp_hook, gcs_hook): gcs_hook.return_value.list.return_value = SOURCE_FILES_LIST[:2] diff --git a/tests/operators/test_local_to_gcs.py b/tests/operators/test_local_to_gcs.py index 921f4ef461374..3f0358cc5e07f 100644 --- a/tests/operators/test_local_to_gcs.py +++ b/tests/operators/test_local_to_gcs.py @@ -55,7 +55,7 @@ def test_init(self): self.assertEqual(operator.mime_type, self._config['mime_type']) self.assertEqual(operator.gzip, self._config['gzip']) - @mock.patch('airflow.operators.local_to_gcs.GoogleCloudStorageHook', + @mock.patch('airflow.operators.local_to_gcs.GcsHook', autospec=True) def test_execute(self, mock_hook): mock_instance = mock_hook.return_value diff --git a/tests/operators/test_mssql_to_gcs.py b/tests/operators/test_mssql_to_gcs.py index bf94fc707d50b..e3ffa74746a33 100644 --- a/tests/operators/test_mssql_to_gcs.py +++ b/tests/operators/test_mssql_to_gcs.py @@ -62,7 +62,7 @@ def test_init(self): self.assertEqual(op.filename, JSON_FILENAME) @mock.patch('airflow.operators.mssql_to_gcs.MsSqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_exec_success_json(self, gcs_hook_mock_class, mssql_hook_mock_class): """Test successful run of execute function for JSON""" op = MsSqlToGoogleCloudStorageOperator( @@ -94,7 +94,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): mssql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL) @mock.patch('airflow.operators.mssql_to_gcs.MsSqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_file_splitting(self, gcs_hook_mock_class, mssql_hook_mock_class): """Test that ndjson is split by approx_max_file_size_bytes param.""" mssql_hook_mock = mssql_hook_mock_class.return_value @@ -125,7 +125,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): op.execute(None) @mock.patch('airflow.operators.mssql_to_gcs.MsSqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_schema_file(self, gcs_hook_mock_class, mssql_hook_mock_class): """Test writing schema files.""" mssql_hook_mock = mssql_hook_mock_class.return_value diff --git a/tests/operators/test_mysql_to_gcs.py b/tests/operators/test_mysql_to_gcs.py index dd4d394d4e40c..8ca5a2100d9a6 100644 --- a/tests/operators/test_mysql_to_gcs.py +++ b/tests/operators/test_mysql_to_gcs.py @@ -101,7 +101,7 @@ def test_convert_type(self, value, schema_type, expected): expected) @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_exec_success_json(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for JSON""" op = MySqlToGoogleCloudStorageOperator( @@ -133,7 +133,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL) @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_exec_success_csv(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for CSV""" op = MySqlToGoogleCloudStorageOperator( @@ -166,7 +166,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL) @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_exec_success_csv_ensure_utc(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for CSV""" op = MySqlToGoogleCloudStorageOperator( @@ -200,7 +200,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): mysql_hook_mock.get_conn().cursor().execute.assert_has_calls([mock.call(TZ_QUERY), mock.call(SQL)]) @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_exec_success_csv_with_delimiter(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test successful run of execute function for CSV with a field delimiter""" op = MySqlToGoogleCloudStorageOperator( @@ -234,7 +234,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): mysql_hook_mock.get_conn().cursor().execute.assert_called_once_with(SQL) @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_file_splitting(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test that ndjson is split by approx_max_file_size_bytes param.""" mysql_hook_mock = mysql_hook_mock_class.return_value @@ -265,7 +265,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type=None, gzip=False): op.execute(None) @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_schema_file(self, gcs_hook_mock_class, mysql_hook_mock_class): """Test writing schema files.""" mysql_hook_mock = mysql_hook_mock_class.return_value @@ -294,7 +294,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip): # pylint: disab self.assertEqual(2, gcs_hook_mock.upload.call_count) @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_query_with_error(self, mock_gcs_hook, mock_mysql_hook): mock_mysql_hook.return_value.get_conn.\ return_value.cursor.return_value.execute.side_effect = ProgrammingError @@ -308,7 +308,7 @@ def test_query_with_error(self, mock_gcs_hook, mock_mysql_hook): op.query() @mock.patch('airflow.operators.mysql_to_gcs.MySqlHook') - @mock.patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @mock.patch('airflow.operators.sql_to_gcs.GcsHook') def test_execute_with_query_error(self, mock_gcs_hook, mock_mysql_hook): mock_mysql_hook.return_value.get_conn.\ return_value.cursor.return_value.execute.side_effect = ProgrammingError diff --git a/tests/operators/test_postgres_to_gcs.py b/tests/operators/test_postgres_to_gcs.py index 56276442e1838..abd6dbde05cc6 100644 --- a/tests/operators/test_postgres_to_gcs.py +++ b/tests/operators/test_postgres_to_gcs.py @@ -82,7 +82,7 @@ def test_init(self): self.assertEqual(op.bucket, BUCKET) self.assertEqual(op.filename, FILENAME) - @patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @patch('airflow.operators.sql_to_gcs.GcsHook') def test_exec_success(self, gcs_hook_mock_class): """Test the execute function in case where the run is successful.""" op = PostgresToGoogleCloudStorageOperator( @@ -106,7 +106,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip): op.execute(None) - @patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @patch('airflow.operators.sql_to_gcs.GcsHook') def test_file_splitting(self, gcs_hook_mock_class): """Test that ndjson is split by approx_max_file_size_bytes param.""" @@ -133,7 +133,7 @@ def _assert_upload(bucket, obj, tmp_filename, mime_type, gzip): approx_max_file_size_bytes=len(expected_upload[FILENAME.format(0)])) op.execute(None) - @patch('airflow.operators.sql_to_gcs.GoogleCloudStorageHook') + @patch('airflow.operators.sql_to_gcs.GcsHook') def test_schema_file(self, gcs_hook_mock_class): """Test writing schema files.""" diff --git a/tests/providers/google/marketing_platform/operators/test_campaign_manager.py b/tests/providers/google/marketing_platform/operators/test_campaign_manager.py index acda7766384ac..8fdaacb69ef34 100644 --- a/tests/providers/google/marketing_platform/operators/test_campaign_manager.py +++ b/tests/providers/google/marketing_platform/operators/test_campaign_manager.py @@ -69,7 +69,7 @@ class TestGoogleCampaignManagerGetReportOperator(TestCase): ) @mock.patch( "airflow.providers.google.marketing_platform.operators." - "campaign_manager.GoogleCloudStorageHook" + "campaign_manager.GcsHook" ) @mock.patch( "airflow.providers.google.marketing_platform.operators." diff --git a/tests/providers/google/marketing_platform/operators/test_display_video.py b/tests/providers/google/marketing_platform/operators/test_display_video.py index 36ade79cbe586..6b4ef7cc56639 100644 --- a/tests/providers/google/marketing_platform/operators/test_display_video.py +++ b/tests/providers/google/marketing_platform/operators/test_display_video.py @@ -96,7 +96,7 @@ class TestGoogleDisplayVideo360GetReportOperator(TestCase): ) @mock.patch( "airflow.providers.google.marketing_platform.operators." - "display_video.GoogleCloudStorageHook" + "display_video.GcsHook" ) @mock.patch( "airflow.providers.google.marketing_platform.operators." diff --git a/tests/providers/google/marketing_platform/operators/test_search_ads.py b/tests/providers/google/marketing_platform/operators/test_search_ads.py index e65d6cfeb3d88..1c45067f52055 100644 --- a/tests/providers/google/marketing_platform/operators/test_search_ads.py +++ b/tests/providers/google/marketing_platform/operators/test_search_ads.py @@ -61,7 +61,7 @@ class TestSearchAdsGetfileReportOperator(TestCase): ) @mock.patch( "airflow.providers.google.marketing_platform." - "operators.search_ads.GoogleCloudStorageHook" + "operators.search_ads.GcsHook" ) @mock.patch( "airflow.providers.google.marketing_platform." diff --git a/tests/test_core_to_contrib.py b/tests/test_core_to_contrib.py index 5f9d07c3c6e56..d39bda4358d40 100644 --- a/tests/test_core_to_contrib.py +++ b/tests/test_core_to_contrib.py @@ -69,7 +69,7 @@ "airflow.contrib.hooks.gcp_text_to_speech_hook.GCPTextToSpeechHook", ), ( - "airflow.gcp.hooks.gcs.GoogleCloudStorageHook", + "airflow.gcp.hooks.gcs.GcsHook", "airflow.contrib.hooks.gcs_hook.GoogleCloudStorageHook", ), (