def test_update_data_source(self):
"""Test that the GCP datasource is updated for given dataset."""
source_id = 13
credentials = {"project_id": "test_project"}
test_dataset_table_id = "test_table_id"
data_source = {"dataset": "test_dataset"}
gcp_source = {
"source_id": source_id,
"source_uuid": FAKE.uuid4(),
"name": "Provider GCP",
"source_type": "GCP",
"authentication": {"credentials": credentials},
"billing_source": {"data_source": data_source},
"auth_header": "fakeauthheader",
"account_id": "acct10001",
"offset": 12,
}
provider = Sources(**gcp_source)
provider.save()
provider = GCPProvider()
updated_data_source = copy.deepcopy(data_source)
updated_data_source["table_id"] = test_dataset_table_id
provider.update_source_data_source(credentials, updated_data_source)
db_obj = Sources.objects.get(source_id=source_id)
self.assertEquals(db_obj.billing_source, {"data_source": updated_data_source})
gcp2_source_id = 14
gcp2_dataset_table_id = "test_table_id_2"
data_source_2 = {"dataset": "test_dataset_2"}
gcp_source2 = {
"source_id": gcp2_source_id,
"source_uuid": FAKE.uuid4(),
"name": "Provider GCP 2",
"source_type": "GCP",
"authentication": {"credentials": credentials},
"billing_source": {"data_source": data_source_2},
"auth_header": "fakeauthheader",
"account_id": "acct10001",
"offset": 12,
}
provider_2 = Sources(**gcp_source2)
provider_2.save()
updated_data_source_2 = copy.deepcopy(data_source_2)
updated_data_source_2["table_id"] = gcp2_dataset_table_id
GCPProvider().update_source_data_source(credentials, updated_data_source_2)
db_obj_2 = Sources.objects.get(source_id=gcp2_source_id)
self.assertEquals(
Sources.objects.get(source_id=source_id).billing_source, {"data_source": updated_data_source}
)
self.assertEquals(db_obj_2.billing_source, {"data_source": updated_data_source_2})
def test_cost_usage_source_is_reachable_valid(self, mock_storage):
"""Test that cost_usage_source_is_reachable succeeds."""
storage_resource_name = {'bucket': FAKE.word()}
credentials = {'project_id': FAKE.word()}
provider = GCPProvider()
self.assertTrue(
provider.cost_usage_source_is_reachable(credentials,
storage_resource_name))
def test_cost_usage_source_raise_google_cloud_error(self, mock_auth, mock_discovery):
"""Test that cost_usage_source_is_reachable succeeds."""
err_msg = "GCP Error"
mock_discovery.build.side_effect = GoogleCloudError(err_msg)
mock_auth.return_value = (MagicMock(), MagicMock())
billing_source_param = {"dataset": FAKE.word(), "table_id": FAKE.word()}
credentials_param = {"project_id": FAKE.word()}
with self.assertRaisesRegexp(ValidationError, err_msg):
provider = GCPProvider()
provider.cost_usage_source_is_reachable(credentials_param, billing_source_param)
def test_cost_usage_source_missing_required_permission(self, mock_auth, mock_discovery):
"""Test that cost_usage_source_is_reachable succeeds."""
missing_permissions = [REQUIRED_IAM_PERMISSIONS[0]]
mock_auth.return_value = (MagicMock(), MagicMock())
mock_discovery.build.return_value.projects.return_value.testIamPermissions.return_value.execute.return_value.get.return_value = ( # noqa: E501
missing_permissions
)
billing_source_param = {"dataset": FAKE.word(), "table_id": FAKE.word()}
credentials_param = {"project_id": FAKE.word()}
with self.assertRaises(ValidationError):
provider = GCPProvider()
provider.cost_usage_source_is_reachable(credentials_param, billing_source_param)
def test_cost_usage_source_is_reachable_dataset_bad_format(self, mock_auth, mock_discovery, mock_bigquery):
"""Test that cost_usage_source_is_reachable throws appropriate error when dataset not correct."""
mock_bigquery.Client.side_effect = BadRequest(message="Incorrect dataset format")
gcp_creds = MagicMock()
mock_auth.return_value = (gcp_creds, MagicMock())
mock_discovery.build.return_value.projects.return_value.testIamPermissions.return_value.execute.return_value.get.return_value = ( # noqa: E501
REQUIRED_IAM_PERMISSIONS
)
billing_source_param = {"dataset": FAKE.word()}
credentials_param = {"project_id": FAKE.word()}
provider = GCPProvider()
with self.assertRaises(ValidationError):
provider.cost_usage_source_is_reachable(credentials_param, billing_source_param)
def test_cost_usage_source_is_reachable_no_table_id_notready(self, mock_auth, mock_discovery, mock_bigquery):
"""Test that cost_usage_source_is_reachable succeeds."""
mock_bigquery.Client.return_value = MockBigQueryClient("test_project", "test_dataset", ["abc_1234"])
gcp_creds = MagicMock()
mock_auth.return_value = (gcp_creds, MagicMock())
mock_discovery.build.return_value.projects.return_value.testIamPermissions.return_value.execute.return_value.get.return_value = ( # noqa: E501
REQUIRED_IAM_PERMISSIONS
)
billing_source_param = {"dataset": FAKE.word()}
credentials_param = {"project_id": FAKE.word()}
provider = GCPProvider()
with self.assertRaises(SkipStatusPush):
provider.cost_usage_source_is_reachable(credentials_param, billing_source_param)
def __init__(self, customer_name, billing_source, **kwargs):
"""
Constructor.
Args:
customer_name (str): Name of the customer
billing_source (dict): dict containing name of GCP storage bucket
"""
super().__init__(**kwargs)
self.bucket_name = billing_source['bucket']
self.customer_name = customer_name.replace(' ', '_')
self._provider_id = kwargs.get('provider_id')
try:
GCPProvider().cost_usage_source_is_reachable(None, billing_source)
self._storage_client = storage.Client()
self._bucket_info = self._storage_client.lookup_bucket(
self.bucket_name)
except ValidationError as ex:
LOG.error(
'GCP bucket %(bucket_name)s for customer %(customer_name)s is not reachable. '
'Error: %(ex)s',
{
'customer_name': customer_name,
'bucket_name': self.bucket_name,
'ex': str(ex),
},
)
raise GCPReportDownloaderError(str(ex))
def __init__(self, customer_name, data_source, **kwargs):
"""
Constructor.
Args:
customer_name (str): Name of the customer
data_source (dict): dict containing name of GCP storage bucket
"""
super().__init__(**kwargs)
self.bucket_name = data_source.get("bucket")
self.report_prefix = data_source.get("report_prefix", "")
self.customer_name = customer_name.replace(" ", "_")
self._provider_uuid = kwargs.get("provider_uuid")
try:
GCPProvider().cost_usage_source_is_reachable(None, data_source)
self._storage_client = storage.Client()
self._bucket_info = self._storage_client.lookup_bucket(
self.bucket_name)
except ValidationError as ex:
msg = f"GCP bucket {self.bucket_name} for customer {customer_name} is not reachable. Error: {str(ex)}"
LOG.error(log_json(self.request_id, msg, self.context))
raise GCPReportDownloaderError(str(ex))
def __init__(self, customer_name, data_source, **kwargs):
"""
Constructor.
Args:
customer_name (str): Name of the customer
data_source (dict): dict containing name of GCP storage bucket
"""
super().__init__(**kwargs)
self.customer_name = customer_name.replace(" ", "_")
self.credentials = kwargs.get("credentials", {})
self.data_source = data_source
self._provider_uuid = kwargs.get("provider_uuid")
self.gcp_big_query_columns = [
"billing_account_id",
"service.id",
"service.description",
"sku.id",
"sku.description",
"usage_start_time",
"usage_end_time",
"project.id",
"project.name",
"project.labels",
"project.ancestry_numbers",
"labels",
"system_labels",
"location.location",
"location.country",
"location.region",
"location.zone",
"export_time",
"cost",
"currency",
"currency_conversion_rate",
"usage.amount",
"usage.unit",
"usage.amount_in_pricing_units",
"usage.pricing_unit",
"credits",
"invoice.month",
"cost_type",
]
self.table_name = ".".join([
self.credentials.get("project_id"),
self._get_dataset_name(),
self.data_source.get("table_id")
])
self.scan_start, self.scan_end = self._generate_default_scan_range()
try:
GCPProvider().cost_usage_source_is_reachable(
self.credentials, self.data_source)
self.etag = self._generate_etag()
except ValidationError as ex:
msg = f"GCP source ({self._provider_uuid}) for {customer_name} is not reachable. Error: {str(ex)}"
LOG.warning(log_json(self.tracing_id, msg, self.context))
raise GCPReportDownloaderError(str(ex))
self.big_query_export_time = None
def __init__(self, task, customer_name, billing_source, **kwargs):
"""
Constructor.
Args:
task (Object) bound celery object
customer_name (str): Name of the customer
billing_source (dict): dict containing name of GCP storage bucket
"""
super().__init__(task, **kwargs)
self.bucket_name = billing_source["bucket"]
self.report_prefix = billing_source.get("report_prefix", "")
self.customer_name = customer_name.replace(" ", "_")
self._provider_uuid = kwargs.get("provider_uuid")
try:
GCPProvider().cost_usage_source_is_reachable(None, billing_source)
self._storage_client = storage.Client()
self._bucket_info = self._storage_client.lookup_bucket(
self.bucket_name)
except ValidationError as ex:
LOG.error(
"GCP bucket %(bucket_name)s for customer %(customer_name)s is not reachable. "
"Error: %(ex)s",
{
"customer_name": customer_name,
"bucket_name": self.bucket_name,
"ex": str(ex)
},
)
raise GCPReportDownloaderError(str(ex))
def test_bucket_access_exception(self, mock_storage):
"""Test that ValidationError is raised when GoogleCloudError is raised."""
gcp_client = mock_storage.return_value
gcp_client.lookup_bucket.side_effect = GoogleCloudError('GCP Error')
credentials = {'project_id': FAKE.word()}
storage_resource_name = {'bucket': FAKE.word()}
with self.assertRaises(ValidationError):
GCPProvider().cost_usage_source_is_reachable(credentials, storage_resource_name)
def test_cost_usage_source_is_reachable_valid(self, mock_auth, mock_discovery, mock_bigquery):
"""Test that cost_usage_source_is_reachable succeeds."""
mock_bigquery.Client.return_value = MockBigQueryClient(
"test_project", "test_dataset", ["gcp_billing_export_1234"]
)
gcp_creds = MagicMock()
mock_auth.return_value = (gcp_creds, MagicMock())
mock_discovery.build.return_value.projects.return_value.testIamPermissions.return_value.execute.return_value.get.return_value = ( # noqa: E501
REQUIRED_IAM_PERMISSIONS
)
billing_source_param = {"dataset": FAKE.word()}
credentials_param = {"project_id": FAKE.word()}
provider = GCPProvider()
self.assertTrue(provider.cost_usage_source_is_reachable(credentials_param, billing_source_param))
mock_discovery.build.assert_called_once_with(
"cloudresourcemanager", "v1", credentials=gcp_creds, cache_discovery=False
)
mock_discovery.build.return_value.projects.assert_called_once()
mock_discovery.build.return_value.projects.return_value.testIamPermissions.assert_called_once()
def test_no_bucket_exists_exception(self, mock_storage):
"""Test that ValidationError is raised when GoogleCloudError is raised."""
gcp_client = mock_storage.return_value
gcp_client.lookup_bucket.return_value = None
credentials = {"project_id": FAKE.word()}
storage_resource_name = {"bucket": FAKE.word()}
with self.assertRaises(ValidationError):
GCPProvider().cost_usage_source_is_reachable(
credentials, storage_resource_name)
def test_format_dataset_id(self):
"""Test helper method _format_dataset_id."""
test_matrix = [
{"dataset": "testdataset", "project_id": "testproject", "expected": "testproject.testdataset"},
{"dataset": "testproject:testdataset", "project_id": "testproject", "expected": "testproject.testdataset"},
]
for test in test_matrix:
data_source = {"dataset": test.get("dataset")}
credentials = {"project_id": test.get("project_id")}
formatted_data_set = GCPProvider()._format_dataset_id(data_source, credentials)
self.assertEqual(formatted_data_set, test.get("expected"))
def test_name(self):
"""Test name property."""
provider = GCPProvider()
self.assertEqual(provider.name(), Provider.PROVIDER_GCP)
def test_name(self):
"""Test name property."""
provider = GCPProvider()
self.assertEqual(provider.name(), 'GCP')
