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"])