def execute(self, context: dict) -> None: hook = GoogleCampaignManagerHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, api_version=self.api_version, impersonation_chain=self.impersonation_chain, ) gcs_hook = GCSHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, impersonation_chain=self.impersonation_chain, ) # Get name of the report report = hook.get_report(file_id=self.file_id, profile_id=self.profile_id, report_id=self.report_id) report_name = self.report_name or report.get("fileName", str(uuid.uuid4())) report_name = self._resolve_file_name(report_name) # Download the report self.log.info("Starting downloading report %s", self.report_id) request = hook.get_report_file(profile_id=self.profile_id, report_id=self.report_id, file_id=self.file_id) with tempfile.NamedTemporaryFile() as temp_file: downloader = http.MediaIoBaseDownload(fd=temp_file, request=request, chunksize=self.chunk_size) download_finished = False while not download_finished: _, download_finished = downloader.next_chunk() temp_file.flush() # Upload the local file to bucket gcs_hook.upload( bucket_name=self.bucket_name, object_name=report_name, gzip=self.gzip, filename=temp_file.name, mime_type="text/csv", ) self.xcom_push(context, key="report_name", value=report_name)
class TestGoogleCampaignManagerHook(TestCase): def setUp(self): with mock.patch( "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__", new=mock_base_gcp_hook_default_project_id, ): self.hook = GoogleCampaignManagerHook(gcp_conn_id=GCP_CONN_ID, api_version=API_VERSION) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook._authorize") @mock.patch( "airflow.providers.google.marketing_platform.hooks.campaign_manager.build" ) def test_gen_conn(self, mock_build, mock_authorize): result = self.hook.get_conn() mock_build.assert_called_once_with( "dfareporting", API_VERSION, http=mock_authorize.return_value, cache_discovery=False, ) assert mock_build.return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_delete_report(self, get_conn_mock): return_value = "TEST" get_conn_mock.return_value.reports.return_value.delete.return_value.execute.return_value = ( return_value) result = self.hook.delete_report(profile_id=PROFILE_ID, report_id=REPORT_ID) get_conn_mock.return_value.reports.return_value.delete.assert_called_once_with( profileId=PROFILE_ID, reportId=REPORT_ID) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_get_report(self, get_conn_mock): file_id = "FILE_ID" return_value = "TEST" # fmt: off get_conn_mock.return_value.reports.return_value.files.return_value. \ get.return_value.execute.return_value = return_value # fmt: on result = self.hook.get_report(profile_id=PROFILE_ID, report_id=REPORT_ID, file_id=file_id) get_conn_mock.return_value.reports.return_value.files.return_value.get.assert_called_once_with( profileId=PROFILE_ID, reportId=REPORT_ID, fileId=file_id) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_get_report_file(self, get_conn_mock): file_id = "FILE_ID" return_value = "TEST" get_conn_mock.return_value.reports.return_value.files.return_value.get_media.return_value = ( return_value) result = self.hook.get_report_file(profile_id=PROFILE_ID, report_id=REPORT_ID, file_id=file_id) get_conn_mock.return_value.reports.return_value.files.return_value.get_media.assert_called_once_with( profileId=PROFILE_ID, reportId=REPORT_ID, fileId=file_id) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_insert_report(self, get_conn_mock): report = {"body": "test"} return_value = "TEST" get_conn_mock.return_value.reports.return_value.insert.return_value.execute.return_value = ( return_value) result = self.hook.insert_report(profile_id=PROFILE_ID, report=report) get_conn_mock.return_value.reports.return_value.insert.assert_called_once_with( profileId=PROFILE_ID, body=report) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_list_reports(self, get_conn_mock): max_results = 42 scope = "SCOPE" sort_field = "SORT_FIELD" sort_order = "SORT_ORDER" items = ["item"] return_value = {"nextPageToken": None, "items": items} get_conn_mock.return_value.reports.return_value.list.return_value.execute.return_value = return_value request_mock = mock.MagicMock() request_mock.execute.return_value = { "nextPageToken": None, "items": items } get_conn_mock.return_value.reports.return_value.list_next.side_effect = [ request_mock, request_mock, request_mock, None, ] result = self.hook.list_reports( profile_id=PROFILE_ID, max_results=max_results, scope=scope, sort_field=sort_field, sort_order=sort_order, ) get_conn_mock.return_value.reports.return_value.list.assert_called_once_with( profileId=PROFILE_ID, maxResults=max_results, scope=scope, sortField=sort_field, sortOrder=sort_order, ) assert items * 4 == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_patch_report(self, get_conn_mock): update_mask = {"test": "test"} return_value = "TEST" get_conn_mock.return_value.reports.return_value.patch.return_value.execute.return_value = return_value result = self.hook.patch_report(profile_id=PROFILE_ID, report_id=REPORT_ID, update_mask=update_mask) get_conn_mock.return_value.reports.return_value.patch.assert_called_once_with( profileId=PROFILE_ID, reportId=REPORT_ID, body=update_mask) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_run_report(self, get_conn_mock): synchronous = True return_value = "TEST" get_conn_mock.return_value.reports.return_value.run.return_value.execute.return_value = return_value result = self.hook.run_report(profile_id=PROFILE_ID, report_id=REPORT_ID, synchronous=synchronous) get_conn_mock.return_value.reports.return_value.run.assert_called_once_with( profileId=PROFILE_ID, reportId=REPORT_ID, synchronous=synchronous) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") def test_update_report(self, get_conn_mock): return_value = "TEST" get_conn_mock.return_value.reports.return_value.update.return_value.execute.return_value = ( return_value) result = self.hook.update_report(profile_id=PROFILE_ID, report_id=REPORT_ID) get_conn_mock.return_value.reports.return_value.update.assert_called_once_with( profileId=PROFILE_ID, reportId=REPORT_ID) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform." "hooks.campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch( "airflow.providers.google.marketing_platform.hooks.campaign_manager.GoogleCampaignManagerHook" "._conversions_batch_request") def test_conversion_batch_insert(self, batch_request_mock, get_conn_mock): conversions = [{"conversions1": "value"}, {"conversions2": "value"}] return_value = {'hasFailures': False} get_conn_mock.return_value.conversions.return_value.batchinsert.return_value.execute.return_value = ( return_value) batch_request_mock.return_value = "batch_request_mock" result = self.hook.conversions_batch_insert( profile_id=PROFILE_ID, conversions=conversions, encryption_entity_id=ENCRYPTION_ENTITY_ID, encryption_entity_type=ENCRYPTION_ENTITY_TYPE, encryption_source=ENCRYPTION_SOURCE, ) batch_request_mock.assert_called_once_with( conversions=conversions, encryption_entity_id=ENCRYPTION_ENTITY_ID, encryption_entity_type=ENCRYPTION_ENTITY_TYPE, encryption_source=ENCRYPTION_SOURCE, kind="dfareporting#conversionsBatchInsertRequest", ) get_conn_mock.return_value.conversions.return_value.batchinsert.assert_called_once_with( profileId=PROFILE_ID, body=batch_request_mock.return_value) assert return_value == result @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch( "airflow.providers.google.marketing_platform.hooks.campaign_manager.GoogleCampaignManagerHook" "._conversions_batch_request") def test_conversions_batch_update(self, batch_request_mock, get_conn_mock): conversions = [{"conversions1": "value"}, {"conversions2": "value"}] return_value = {'hasFailures': False} get_conn_mock.return_value.conversions.return_value.batchupdate.return_value.execute.return_value = ( return_value) batch_request_mock.return_value = "batch_request_mock" result = self.hook.conversions_batch_update( profile_id=PROFILE_ID, conversions=conversions, encryption_entity_id=ENCRYPTION_ENTITY_ID, encryption_entity_type=ENCRYPTION_ENTITY_TYPE, encryption_source=ENCRYPTION_SOURCE, ) batch_request_mock.assert_called_once_with( conversions=conversions, encryption_entity_id=ENCRYPTION_ENTITY_ID, encryption_entity_type=ENCRYPTION_ENTITY_TYPE, encryption_source=ENCRYPTION_SOURCE, kind="dfareporting#conversionsBatchUpdateRequest", ) get_conn_mock.return_value.conversions.return_value.batchupdate.assert_called_once_with( profileId=PROFILE_ID, body=batch_request_mock.return_value) assert return_value == result
class TestGoogleCampaignManagerHook(TestCase): def setUp(self): with mock.patch( "airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__", new=mock_base_gcp_hook_default_project_id, ): self.hook = GoogleCampaignManagerHook(gcp_conn_id=GCP_CONN_ID, api_version=API_VERSION) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook._authorize") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.build") def test_gen_conn(self, mock_build, mock_authorize): result = self.hook.get_conn() mock_build.assert_called_once_with( "dfareporting", API_VERSION, http=mock_authorize.return_value, cache_discovery=False, ) self.assertEqual(mock_build.return_value, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_delete_report(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" report_id = "REPORT_ID" return_value = "TEST" get_conn_mock.return_value.reports.return_value.delete.return_value.execute.return_value = ( return_value) result = self.hook.delete_report(profile_id=profile_id, report_id=report_id) get_conn_mock.return_value.reports.return_value.delete.assert_called_once_with( profileId=profile_id, reportId=report_id) self.assertEqual(return_value, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_get_report(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" report_id = "REPORT_ID" file_id = "FILE_ID" return_value = "TEST" get_conn_mock.return_value.reports.return_value.files.return_value.get.\ return_value.execute.return_value = return_value result = self.hook.get_report(profile_id=profile_id, report_id=report_id, file_id=file_id) get_conn_mock.return_value.reports.return_value.files.return_value.get.assert_called_once_with( profileId=profile_id, reportId=report_id, fileId=file_id) self.assertEqual(return_value, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_get_report_file(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" report_id = "REPORT_ID" file_id = "FILE_ID" return_value = "TEST" get_conn_mock.return_value.reports.return_value.files.return_value.get_media.return_value = ( return_value) result = self.hook.get_report_file(profile_id=profile_id, report_id=report_id, file_id=file_id) get_conn_mock.return_value.reports.return_value.files.return_value.get_media.assert_called_once_with( profileId=profile_id, reportId=report_id, fileId=file_id) self.assertEqual(return_value, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_insert_report(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" report = {"body": "test"} return_value = "TEST" get_conn_mock.return_value.reports.return_value.insert.return_value.execute.return_value = ( return_value) result = self.hook.insert_report(profile_id=profile_id, report=report) get_conn_mock.return_value.reports.return_value.insert.assert_called_once_with( profileId=profile_id, body=report) self.assertEqual(return_value, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_list_reports(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" max_results = 42 scope = "SCOPE" sort_field = "SORT_FIELD" sort_order = "SORT_ORDER" items = ['item'] return_value = {"nextPageToken": None, "items": items} get_conn_mock.return_value.reports.return_value.list.return_value.\ execute.return_value = return_value request_mock = mock.MagicMock() request_mock.execute.return_value = { "nextPageToken": None, "items": items } get_conn_mock.return_value.reports.return_value.list_next.side_effect = [ request_mock, request_mock, request_mock, None ] result = self.hook.list_reports( profile_id=profile_id, max_results=max_results, scope=scope, sort_field=sort_field, sort_order=sort_order, ) get_conn_mock.return_value.reports.return_value.list.assert_called_once_with( profileId=profile_id, maxResults=max_results, scope=scope, sortField=sort_field, sortOrder=sort_order, ) self.assertEqual(items * 4, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_patch_report(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" report_id = "REPORT_ID" update_mask = {"test": "test"} return_value = "TEST" get_conn_mock.return_value.reports.return_value.patch.return_value.execute.return_value = ( return_value) result = self.hook.patch_report(profile_id=profile_id, report_id=report_id, update_mask=update_mask) get_conn_mock.return_value.reports.return_value.patch.assert_called_once_with( profileId=profile_id, reportId=report_id, body=update_mask) self.assertEqual(return_value, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_run_report(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" report_id = "REPORT_ID" synchronous = True return_value = "TEST" get_conn_mock.return_value.reports.return_value.run.return_value.execute.return_value = ( return_value) result = self.hook.run_report(profile_id=profile_id, report_id=report_id, synchronous=synchronous) get_conn_mock.return_value.reports.return_value.run.assert_called_once_with( profileId=profile_id, reportId=report_id, synchronous=synchronous) self.assertEqual(return_value, result) @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.GoogleCampaignManagerHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "campaign_manager.CloudBaseHook.__init__") def test_update_report(self, mock_base_hook, get_conn_mock): profile_id = "PROFILE_ID" report_id = "REPORT_ID" return_value = "TEST" get_conn_mock.return_value.reports.return_value.update.return_value.execute.return_value = ( return_value) result = self.hook.update_report(profile_id=profile_id, report_id=report_id) get_conn_mock.return_value.reports.return_value.update.assert_called_once_with( profileId=profile_id, reportId=report_id) self.assertEqual(return_value, result)