def test_execute(self, mock_temp, gcs_mock_hook, mock_hook):
operation_name = "operation_name"
operation = {"key": "value"}
bucket_name = "bucket_name"
object_name = "object_name"
filename = "filename"
gzip = False
# mock_hook.return_value.create_sdf_download_operation.return_value = response_name
mock_hook.return_value.get_sdf_download_operation.return_value = operation
mock_temp.NamedTemporaryFile.return_value.__enter__.return_value.name = filename
op = GoogleDisplayVideo360SDFtoGCSOperator(
operation_name=operation_name,
bucket_name=bucket_name,
object_name=object_name,
gzip=gzip,
api_version=API_VERSION,
gcp_conn_id=GCP_CONN_ID,
task_id="test_task",
)
op.execute(context=None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID, api_version=API_VERSION, delegate_to=DELEGATE_TO
)
mock_hook.return_value.get_sdf_download_operation.assert_called_once()
mock_hook.return_value.get_sdf_download_operation.assert_called_once_with(
operation_name=operation_name
)
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
)
mock_hook.return_value.download_content_from_request.assert_called_once()
mock_hook.return_value.download_content_from_request.assert_called_once_with(
mock_temp.NamedTemporaryFile.return_value.__enter__.return_value,
mock_hook.return_value.download_media(),
chunk_size=1024 * 1024,
)
gcs_mock_hook.assert_called_once()
gcs_mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID, delegate_to=DELEGATE_TO
)
gcs_mock_hook.return_value.upload.assert_called_once()
gcs_mock_hook.return_value.upload.assert_called_once_with(
bucket_name=bucket_name,
object_name=object_name,
filename=filename,
gzip=gzip,
)
)
operation_name = '{{ task_instance.xcom_pull("create_sdf_download_task")["name"] }}'
# [END howto_google_display_video_create_sdf_download_task_operator]
# [START howto_google_display_video_wait_for_operation_sensor]
wait_for_operation = GoogleDisplayVideo360GetSDFDownloadOperationSensor(
task_id="wait_for_operation",
operation_name=operation_name,
)
# [END howto_google_display_video_wait_for_operation_sensor]
# [START howto_google_display_video_save_sdf_in_gcs_operator]
save_sdf_in_gcs = GoogleDisplayVideo360SDFtoGCSOperator(
task_id="save_sdf_in_gcs",
operation_name=operation_name,
bucket_name=BUCKET,
object_name=BUCKET_FILE_LOCATION,
gzip=False,
)
# [END howto_google_display_video_save_sdf_in_gcs_operator]
# [START howto_google_display_video_gcs_to_big_query_operator]
upload_sdf_to_big_query = GCSToBigQueryOperator(
task_id="upload_sdf_to_big_query",
bucket=BUCKET,
source_objects=['{{ task_instance.xcom_pull("upload_sdf_to_bigquery")}}'],
destination_project_dataset_table=f"{BQ_DATA_SET}.gcs_to_bq_table",
schema_fields=[
{"name": "name", "type": "STRING", "mode": "NULLABLE"},
{"name": "post_abbr", "type": "STRING", "mode": "NULLABLE"},
],