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: