def test_execute(
self,
mock_base_op,
mock_hook,
mock_gcs_hook,
mock_xcom,
mock_temp,
mock_reuqest,
mock_shutil,
):
report_id = "REPORT_ID"
bucket_name = "BUCKET"
report_name = "TEST.csv"
filename = "test"
mock_temp.NamedTemporaryFile.return_value.__enter__.return_value.name = filename
mock_hook.return_value.get_query.return_value = {
"metadata": {
"running": False,
"googleCloudStoragePathForLatestReport": "test",
}
}
op = GoogleDisplayVideo360DownloadReportOperator(
report_id=report_id,
api_version=API_VERSION,
bucket_name=bucket_name,
report_name=report_name,
task_id="test_task",
)
op.execute(context=None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
delegate_to=None,
api_version=API_VERSION,
impersonation_chain=None,
)
mock_hook.return_value.get_query.assert_called_once_with(
query_id=report_id)
mock_gcs_hook.assert_called_once_with(
google_cloud_storage_conn_id=GCP_CONN_ID,
delegate_to=None,
impersonation_chain=None,
)
mock_gcs_hook.return_value.upload.assert_called_once_with(
bucket_name=bucket_name,
filename=filename,
gzip=True,
mime_type="text/csv",
object_name=report_name + ".gz",
)
mock_xcom.assert_called_once_with(None,
key="report_name",
value=report_name + ".gz")
# [END howto_google_display_video_createquery_report_operator]
# [START howto_google_display_video_runquery_report_operator]
run_report = GoogleDisplayVideo360RunReportOperator(
report_id=report_id, params=PARAMS, task_id="run_report"
)
# [END howto_google_display_video_runquery_report_operator]
# [START howto_google_display_video_wait_report_operator]
wait_for_report = GoogleDisplayVideo360ReportSensor(task_id="wait_for_report", report_id=report_id)
# [END howto_google_display_video_wait_report_operator]
# [START howto_google_display_video_getquery_report_operator]
get_report = GoogleDisplayVideo360DownloadReportOperator(
report_id=report_id,
task_id="get_report",
bucket_name=BUCKET,
report_name="test1.csv",
)
# [END howto_google_display_video_getquery_report_operator]
# [START howto_google_display_video_deletequery_report_operator]
delete_report = GoogleDisplayVideo360DeleteReportOperator(report_id=report_id, task_id="delete_report")
# [END howto_google_display_video_deletequery_report_operator]
create_report >> run_report >> wait_for_report >> get_report >> delete_report
with models.DAG(
"example_display_video_misc",
schedule_interval=None, # Override to match your needs,
start_date=dates.days_ago(1),
) as dag2:
