Fix download_media url in GoogleDisplayVideo360SDFtoGCSOperator (#…

…22479)
apache · Mar 23, 2022 · 38fde2e · 38fde2e
1 parent 0f0a1a7
commit 38fde2e
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/airflow/providers/google/marketing_platform/hooks/display_video.py b/airflow/providers/google/marketing_platform/hooks/display_video.py
@@ -27,7 +27,7 @@
 class GoogleDisplayVideo360Hook(GoogleBaseHook):
  """Hook for Google Display & Video 360."""
 
- _conn = None # type: Optional[Any]
+ _conn: Optional[Resource] = None
 
  def __init__(
  self,

diff --git a/airflow/providers/google/marketing_platform/operators/display_video.py b/airflow/providers/google/marketing_platform/operators/display_video.py
@@ -692,7 +692,7 @@ def execute(self, context: 'Context') -> str:
  operation_state = hook.get_sdf_download_operation(operation_name=self.operation_name)
 
  self.log.info("Creating file for upload...")
- media = hook.download_media(resource_name=operation_state)
+ media = hook.download_media(resource_name=operation_state["response"]["resourceName"])
 
  self.log.info("Sending file to the Google Cloud Storage...")
  with tempfile.NamedTemporaryFile() as temp_file:

diff --git a/tests/providers/google/marketing_platform/operators/test_display_video.py b/tests/providers/google/marketing_platform/operators/test_display_video.py
@@ -354,7 +354,6 @@ class TestGoogleDisplayVideo360SDFtoGCSOperator(TestCase):
  @mock.patch("airflow.providers.google.marketing_platform.operators.display_video.tempfile")
  def test_execute(self, mock_temp, gcs_mock_hook, mock_hook):
  operation_name = "operation_name"
- operation = {"key": "value"}
  operation = {"response": {"resourceName": "test_name"}}
  gzip = False
 
@@ -388,7 +387,9 @@ def test_execute(self, mock_temp, gcs_mock_hook, mock_hook):
 
  mock_hook.return_value.download_media.assert_called_once()
  mock_hook.return_value.download_media.assert_called_once_with(
- resource_name=mock_hook.return_value.get_sdf_download_operation.return_value
+ resource_name=mock_hook.return_value.get_sdf_download_operation.return_value["response"][
+ "resourceName"
+ ]
  )
 
  mock_hook.return_value.download_content_from_request.assert_called_once()