Remove deprecated method call (blob.download_as_string) (#20091)

kazanzhy · web-flow · commit 50bf53665649 · 2021-12-08T00:53:15.000+01:00
diff --git a/airflow/providers/google/cloud/hooks/gcs.py b/airflow/providers/google/cloud/hooks/gcs.py
@@ -311,7 +311,7 @@ def download(
                     self.log.info('File downloaded to %s', filename)
                     return filename
                 else:
-                    return blob.download_as_string()
+                    return blob.download_as_bytes()
 
             except GoogleCloudError:
                 if num_file_attempts == num_max_attempts:
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
@@ -948,7 +948,7 @@ def execute(self, context) -> None:
                 delegate_to=self.delegate_to,
                 impersonation_chain=self.impersonation_chain,
             )
-            schema_fields = json.loads(gcs_hook.download(gcs_bucket, gcs_object))
+            schema_fields = json.loads(gcs_hook.download(gcs_bucket, gcs_object).decode("utf-8"))
         else:
             schema_fields = self.schema_fields
 
@@ -1203,7 +1203,7 @@ def execute(self, context) -> None:
                 delegate_to=self.delegate_to,
                 impersonation_chain=self.impersonation_chain,
             )
-            schema_fields = json.loads(gcs_hook.download(self.bucket, self.schema_object))
+            schema_fields = json.loads(gcs_hook.download(self.bucket, self.schema_object).decode("utf-8"))
         else:
             schema_fields = self.schema_fields
 
diff --git a/airflow/providers/google/cloud/utils/mlengine_operator_utils.py b/airflow/providers/google/cloud/utils/mlengine_operator_utils.py
@@ -264,7 +264,7 @@ def apply_validate_fn(*args, templates_dict, **kwargs):
             raise ValueError(f"Wrong format prediction_path: {prediction_path}")
         summary = os.path.join(obj.strip("/"), "prediction.summary.json")
         gcs_hook = GCSHook()
-        summary = json.loads(gcs_hook.download(bucket, summary))
+        summary = json.loads(gcs_hook.download(bucket, summary).decode("utf-8"))
         return validate_fn(summary)
 
     evaluate_validation = PythonOperator(
diff --git a/tests/providers/google/cloud/hooks/test_gcs.py b/tests/providers/google/cloud/hooks/test_gcs.py
@@ -641,12 +641,12 @@ def test_compose_without_destination_object(self, mock_service):
         assert str(ctx.value) == 'bucket_name and destination_object cannot be empty.'
 
     @mock.patch(GCS_STRING.format('GCSHook.get_conn'))
-    def test_download_as_string(self, mock_service):
+    def test_download_as_bytes(self, mock_service):
         test_bucket = 'test_bucket'
         test_object = 'test_object'
         test_object_bytes = io.BytesIO(b"input")
 
-        download_method = mock_service.return_value.bucket.return_value.blob.return_value.download_as_string
+        download_method = mock_service.return_value.bucket.return_value.blob.return_value.download_as_bytes
         download_method.return_value = test_object_bytes
 
         response = self.gcs_hook.download(bucket_name=test_bucket, object_name=test_object, filename=None)
@@ -666,10 +666,10 @@ def test_download_to_file(self, mock_service):
         )
         download_filename_method.return_value = None
 
-        download_as_a_string_method = (
-            mock_service.return_value.bucket.return_value.blob.return_value.download_as_string
+        download_as_a_bytes_method = (
+            mock_service.return_value.bucket.return_value.blob.return_value.download_as_bytes
         )
-        download_as_a_string_method.return_value = test_object_bytes
+        download_as_a_bytes_method.return_value = test_object_bytes
         response = self.gcs_hook.download(
             bucket_name=test_bucket, object_name=test_object, filename=test_file
         )
@@ -690,10 +690,10 @@ def test_provide_file(self, mock_service, mock_temp_file):
         )
         download_filename_method.return_value = None
 
-        download_as_a_string_method = (
-            mock_service.return_value.bucket.return_value.blob.return_value.download_as_string
+        download_as_a_bytes_method = (
+            mock_service.return_value.bucket.return_value.blob.return_value.download_as_bytes
         )
-        download_as_a_string_method.return_value = test_object_bytes
+        download_as_a_bytes_method.return_value = test_object_bytes
         mock_temp_file.return_value.__enter__.return_value = mock.MagicMock()
         mock_temp_file.return_value.__enter__.return_value.name = test_file
 
diff --git a/tests/providers/google/cloud/operators/test_mlengine_utils.py b/tests/providers/google/cloud/operators/test_mlengine_utils.py
@@ -127,7 +127,7 @@ def test_successful_run(self):
 
         with patch('airflow.providers.google.cloud.utils.mlengine_operator_utils.GCSHook') as mock_gcs_hook:
             hook_instance = mock_gcs_hook.return_value
-            hook_instance.download.return_value = '{"err": 0.9, "count": 9}'
+            hook_instance.download.return_value = b'{"err": 0.9, "count": 9}'
             result = validate.execute({})
             hook_instance.download.assert_called_once_with(
                 'legal-bucket', 'fake-output-path/prediction.summary.json'
diff --git a/tests/providers/google/cloud/utils/test_mlengine_operator_utils.py b/tests/providers/google/cloud/utils/test_mlengine_operator_utils.py
@@ -232,7 +232,7 @@ def test_apply_validate_fn(self, mock_beam_pipeline, mock_python, mock_download)
 
         _, _, evaluate_validation = result
 
-        mock_download.return_value = json.dumps({"err": 0.3, "mse": 0.04, "count": 1100})
+        mock_download.return_value = json.dumps({"err": 0.3, "mse": 0.04, "count": 1100}).encode("utf-8")
         templates_dict = {"prediction_path": PREDICTION_PATH}
         with pytest.raises(ValueError) as ctx:
             evaluate_validation.python_callable(templates_dict=templates_dict)
