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')