def execute(self, context: 'Context') -> None: gcs_hook = GCSHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, impersonation_chain=self.impersonation_chain, ) ga_hook = GoogleAnalyticsHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, api_version=self.api_version, impersonation_chain=self.impersonation_chain, ) with NamedTemporaryFile("w+") as tmp_file: self.log.info( "Downloading file from GCS: %s/%s ", self.storage_bucket, self.storage_name_object, ) gcs_hook.download( bucket_name=self.storage_bucket, object_name=self.storage_name_object, filename=tmp_file.name, ) ga_hook.upload_data( tmp_file.name, self.account_id, self.web_property_id, self.custom_data_source_id, self.resumable_upload, )
class TestGoogleAnalyticsHook(unittest.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 = GoogleAnalyticsHook(API_VERSION, GCP_CONN_ID) @mock.patch("airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__") def test_init(self, mock_base_init): hook = GoogleAnalyticsHook( API_VERSION, GCP_CONN_ID, delegate_to=DELEGATE_TO, impersonation_chain=IMPERSONATION_CHAIN, ) mock_base_init.assert_called_once_with( GCP_CONN_ID, delegate_to=DELEGATE_TO, impersonation_chain=IMPERSONATION_CHAIN, ) self.assertEqual(hook.api_version, API_VERSION) @mock.patch( "airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook._authorize" ) @mock.patch("airflow.providers.google.marketing_platform.hooks.analytics.build") def test_gen_conn(self, mock_build, mock_authorize): result = self.hook.get_conn() mock_build.assert_called_once_with( "analytics", 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." "analytics.GoogleAnalyticsHook.get_conn") def test_list_accounts(self, get_conn_mock): mock_accounts = get_conn_mock.return_value.management.return_value.accounts mock_list = mock_accounts.return_value.list mock_execute = mock_list.return_value.execute mock_execute.return_value = {"items": ["a", "b"], "totalResults": 2} list_accounts = self.hook.list_accounts() self.assertEqual(list_accounts, ["a", "b"]) @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook.get_conn") def test_list_accounts_for_multiple_pages(self, get_conn_mock): mock_accounts = get_conn_mock.return_value.management.return_value.accounts mock_list = mock_accounts.return_value.list mock_execute = mock_list.return_value.execute mock_execute.side_effect = [ {"items": ["a"], "totalResults": 2}, {"items": ["b"], "totalResults": 2}, ] list_accounts = self.hook.list_accounts() self.assertEqual(list_accounts, ["a", "b"]) @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook.get_conn") def test_get_ad_words_links_call(self, get_conn_mock): num_retries = 5 self.hook.get_ad_words_link( account_id=ACCOUNT_ID, web_property_id=WEB_PROPERTY_ID, web_property_ad_words_link_id=WEB_PROPERTY_AD_WORDS_LINK_ID, ) # fmt: off get_conn_mock.return_value.management.return_value.webPropertyAdWordsLinks.\ return_value.get.return_value.execute.assert_called_once_with( num_retries=num_retries ) get_conn_mock.return_value.management.return_value.webPropertyAdWordsLinks.\ return_value.get.assert_called_once_with( accountId=ACCOUNT_ID, webPropertyId=WEB_PROPERTY_ID, webPropertyAdWordsLinkId=WEB_PROPERTY_AD_WORDS_LINK_ID, ) # fmt: on @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook.get_conn") def test_list_ad_words_links(self, get_conn_mock): mock_ads_links = get_conn_mock.return_value.management.return_value.webPropertyAdWordsLinks mock_list = mock_ads_links.return_value.list mock_execute = mock_list.return_value.execute mock_execute.return_value = {"items": ["a", "b"], "totalResults": 2} list_ads_links = self.hook.list_ad_words_links(account_id=ACCOUNT_ID, web_property_id=WEB_PROPERTY_ID) self.assertEqual(list_ads_links, ["a", "b"]) @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook.get_conn") def test_list_ad_words_links_for_multiple_pages(self, get_conn_mock): mock_ads_links = get_conn_mock.return_value.management.return_value.webPropertyAdWordsLinks mock_list = mock_ads_links.return_value.list mock_execute = mock_list.return_value.execute mock_execute.side_effect = [ {"items": ["a"], "totalResults": 2}, {"items": ["b"], "totalResults": 2}, ] list_ads_links = self.hook.list_ad_words_links(account_id=ACCOUNT_ID, web_property_id=WEB_PROPERTY_ID) self.assertEqual(list_ads_links, ["a", "b"]) @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook.get_conn") @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.MediaFileUpload") def test_upload_data(self, media_mock, get_conn_mock): temp_name = "temp/file" self.hook.upload_data( file_location=temp_name, account_id=ACCOUNT_ID, web_property_id=WEB_PROPERTY_ID, custom_data_source_id=DATA_SOURCE, resumable_upload=True, ) media_mock.assert_called_once_with(temp_name, mimetype="application/octet-stream", resumable=True) # fmt: off get_conn_mock.return_value.management.return_value.uploads.return_value.uploadData.\ assert_called_once_with( accountId=ACCOUNT_ID, webPropertyId=WEB_PROPERTY_ID, customDataSourceId=DATA_SOURCE, media_body=media_mock.return_value, ) # fmt: on @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook.get_conn") def test_delete_upload_data(self, get_conn_mock): body = {"key": "temp/file"} self.hook.delete_upload_data( account_id=ACCOUNT_ID, web_property_id=WEB_PROPERTY_ID, custom_data_source_id=DATA_SOURCE, delete_request_body=body, ) # fmt: off get_conn_mock.return_value.management.return_value.uploads.return_value.deleteUploadData.\ assert_called_once_with( accountId=ACCOUNT_ID, webPropertyId=WEB_PROPERTY_ID, customDataSourceId=DATA_SOURCE, body=body, ) # fmt: on @mock.patch("airflow.providers.google.marketing_platform.hooks." "analytics.GoogleAnalyticsHook.get_conn") def test_list_upload(self, get_conn_mock): uploads = get_conn_mock.return_value.management.return_value.uploads.return_value uploads.list.return_value.execute.return_value = { "items": ["a", "b"], "totalResults": 2, } result = self.hook.list_uploads( account_id=ACCOUNT_ID, web_property_id=WEB_PROPERTY_ID, custom_data_source_id=DATA_SOURCE, ) self.assertEqual(result, ["a", "b"])