def test_prepare_template(self): body = {"key": "value"} with NamedTemporaryFile("w+", suffix=".json") as f: f.write(json.dumps(body)) f.flush() op = GoogleDisplayVideo360CreateReportOperator( body=body, api_version=API_VERSION, task_id="test_task") op.prepare_template() assert isinstance(op.body, dict) assert op.body == body
def test_execute(self, mock_base_op, hook_mock, xcom_mock): body = {"body": "test"} query_id = "TEST" hook_mock.return_value.create_query.return_value = {"queryId": query_id} op = GoogleDisplayVideo360CreateReportOperator( body=body, api_version=API_VERSION, task_id="test_task" ) op.execute(context=None) hook_mock.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, delegate_to=None, api_version=API_VERSION ) hook_mock.return_value.create_query.assert_called_once_with(query=body) xcom_mock.assert_called_once_with(None, key="report_id", value=query_id)
CREATE_SDF_DOWNLOAD_TASK_BODY_REQUEST: Dict = { "version": SDF_VERSION, "advertiserId": ADVERTISER_ID, "inventorySourceFilter": {"inventorySourceIds": []}, } DOWNLOAD_LINE_ITEMS_REQUEST: Dict = {"filterType": ADVERTISER_ID, "format": "CSV", "fileSpec": "EWF"} # [END howto_display_video_env_variables] with models.DAG( "example_display_video", schedule_interval=None, # Override to match your needs, start_date=dates.days_ago(1), ) as dag1: # [START howto_google_display_video_createquery_report_operator] create_report = GoogleDisplayVideo360CreateReportOperator(body=REPORT, task_id="create_report") report_id = "{{ task_instance.xcom_pull('create_report', key='report_id') }}" # [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(