def test_create_budget(self): # Setup Expected Response name = "name3373707" display_name = "displayName1615086568" etag = "etag3123477" expected_response = {"name": name, "display_name": display_name, "etag": etag} expected_response = budget_model_pb2.Budget(**expected_response) # Mock the API response channel = ChannelStub(responses=[expected_response]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = billing_budgets_v1beta1.BudgetServiceClient() # Setup Request parent = client.billing_account_path("[BILLING_ACCOUNT]") budget = {} response = client.create_budget(parent, budget) assert expected_response == response assert len(channel.requests) == 1 expected_request = budget_service_pb2.CreateBudgetRequest( parent=parent, budget=budget ) actual_request = channel.requests[0][1] assert expected_request == actual_request
def test_list_budgets(self): # Setup Expected Response next_page_token = "" budgets_element = {} budgets = [budgets_element] expected_response = {"next_page_token": next_page_token, "budgets": budgets} expected_response = budget_service_pb2.ListBudgetsResponse(**expected_response) # Mock the API response channel = ChannelStub(responses=[expected_response]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = billing_budgets_v1beta1.BudgetServiceClient() # Setup Request parent = client.billing_account_path("[BILLING_ACCOUNT]") paged_list_response = client.list_budgets(parent) resources = list(paged_list_response) assert len(resources) == 1 assert expected_response.budgets[0] == resources[0] assert len(channel.requests) == 1 expected_request = budget_service_pb2.ListBudgetsRequest(parent=parent) actual_request = channel.requests[0][1] assert expected_request == actual_request
def test_get_budget(self): # Setup Expected Response name_2 = "name2-1052831874" display_name = "displayName1615086568" etag = "etag3123477" expected_response = {"name": name_2, "display_name": display_name, "etag": etag} expected_response = budget_model_pb2.Budget(**expected_response) # Mock the API response channel = ChannelStub(responses=[expected_response]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = billing_budgets_v1beta1.BudgetServiceClient() # Setup Request name = client.budget_path("[BILLING_ACCOUNT]", "[BUDGET]") response = client.get_budget(name) assert expected_response == response assert len(channel.requests) == 1 expected_request = budget_service_pb2.GetBudgetRequest(name=name) actual_request = channel.requests[0][1] assert expected_request == actual_request
def test_delete_budget_exception(self): # Mock the API response channel = ChannelStub(responses=[CustomException()]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = billing_budgets_v1beta1.BudgetServiceClient() # Setup request name = client.budget_path("[BILLING_ACCOUNT]", "[BUDGET]") with pytest.raises(CustomException): client.delete_budget(name)
def test_list_budgets_exception(self): channel = ChannelStub(responses=[CustomException()]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = billing_budgets_v1beta1.BudgetServiceClient() # Setup request parent = client.billing_account_path("[BILLING_ACCOUNT]") paged_list_response = client.list_budgets(parent) with pytest.raises(CustomException): list(paged_list_response)
def test_update_budget_exception(self): # Mock the API response channel = ChannelStub(responses=[CustomException()]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = billing_budgets_v1beta1.BudgetServiceClient() # Setup request budget = {} with pytest.raises(CustomException): client.update_budget(budget)
def test_delete_budget(self): channel = ChannelStub() patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = billing_budgets_v1beta1.BudgetServiceClient() # Setup Request name = client.budget_path("[BILLING_ACCOUNT]", "[BUDGET]") client.delete_budget(name) assert len(channel.requests) == 1 expected_request = budget_service_pb2.DeleteBudgetRequest(name=name) actual_request = channel.requests[0][1] assert expected_request == actual_request
def main(gcs_bucket, gcs_file_path, billing_account_id, default_pubsub_topic, mysql_host, mysql_port, mysql_user, mysql_pass, mysql_db, statsd_host, local_mode, dry_run): logformat = '%(asctime)s - %(name)s - %(levelname)s - %(message)s' logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=logformat) # TODO - troubleshoot metric uploading with local container, and test it works on GKE as well. # Remove this line to see WARNINGs logging.getLogger('datadog.dogstatsd').setLevel(logging.ERROR) # Load config.yaml file if local_mode: with open("config.yaml", "r") as f: budget_dict = yaml.load(f, Loader=yaml.SafeLoader) setup_metrics("localhost") else: # Load yaml configuration file from GCS gcs_client = storage.Client() bucket = gcs_client.bucket(gcs_bucket) blob = bucket.blob(gcs_file_path) yaml_string = blob.download_as_string() budget_dict = yaml.load(yaml_string, Loader=yaml.SafeLoader) # Setup metrics with configured statsd host setup_metrics(statsd_host) # Setup gcp helper billing_client = billing_budgets_v1beta1.BudgetServiceClient() resource_manager_client = resource_manager.Client() gcp = GcpHelper(billing_client, resource_manager_client) # Setup mysql connection pool pool = PooledDB(creator=pymysql, host=mysql_host, user=mysql_user, password=mysql_pass, database=mysql_db, autocommit=True, blocking=True, maxconnections=5) mysql_conn = pool.connection() # Iterate over all configured project budgets for project_dict in budget_dict['projects']: config_type = PLUTUS_CONFIG_TYPE_PROJECT if verify_project_yaml(project_dict): project = ProjectBudget(project_dict, config_type, billing_account_id, default_pubsub_topic) budget = gcp.get_and_update_or_create_budget(project) if budget is not None: with mysql_conn.cursor() as mysql_cursor: upsert_budget(mysql_cursor, budget, project.project_id, config_type, project.alert_emails) else: log.error("Project config verification failed.") metrics.incr("error_count", tags=[ "type:misconfig", f"project_id:{project.project_id}", f"config_type:{config_type}" ]) sys.exit(1) # Iterate over all configured parent folder id budgets for parent_dict in budget_dict['parent_folders']: parent_id = parent_dict['parent_folder_id'] config_type = PLUTUS_CONFIG_TYPE_PARENT if verify_parent_yaml(parent_dict): parent_filter = {'parent.id': parent_id} for p in resource_manager_client.list_projects(parent_filter): # Overwrite 'project_id' key for each project under parent folder parent_dict['project_id'] = p.project_id project = ProjectBudget(parent_dict, config_type, billing_account_id, default_pubsub_topic) existing_project_budget_id = gcp.has_existing_project_budget( project) if existing_project_budget_id is None: # No project one off budget configured for this projectid budget = gcp.get_and_update_or_create_budget(project) if budget is not None: with mysql_conn.cursor() as mysql_cursor: upsert_budget(mysql_cursor, budget, project.project_id, config_type, project.alert_emails) else: log.info(f"Skipping creating parent project budget for \ {p.project_id} since exiting plutus project budget found." ) existing_parent_budget_id = gcp.has_existing_parent_budget( parent_id, project) if existing_parent_budget_id is not None: # We have a configured plutus project budget. Delete parent budget gcp.delete_budget(existing_parent_budget_id) else: log.error("Parent folder config verification failed.") metrics.incr("error_count", tags=[ "type:misconfig", f"parent_id:{parent_id}", f"config_type:{config_type}" ]) sys.exit(1) # Iterate over all configured label budgets for label_dict in budget_dict['labels']: config_type = PLUTUS_CONFIG_TYPE_LABEL if verify_labels_yaml(label_dict): labels_filter = {} for row in label_dict['label_list']: for key in row: labels_filter[f"labels.{key}"] = row[key] # Find projects that match the labels for p in resource_manager_client.list_projects( filter_params=labels_filter): # Overwrite the 'project_id' key for each project that matches labels label_dict['project_id'] = p.project_id project = ProjectBudget(label_dict, config_type, billing_account_id, default_pubsub_topic) existing_project_budget_id = gcp.has_existing_project_budget( project) if existing_project_budget_id is None: # No project one off budget configured for this projectid budget = gcp.get_and_update_or_create_budget(project) if budget is not None: with mysql_conn.cursor() as mysql_cursor: upsert_budget(mysql_cursor, budget, project.project_id, config_type, project.alert_emails) else: log.info(f"Skipping creating label project budget for \ {p.project_id} since exiting plutus project budget found." ) existing_labels_budget_id = gcp.has_existing_labels_budget( project) if existing_labels_budget_id is not None: # We have a configured plutus project budget. Delete labels budget gcp.delete_budget(existing_labels_budget_id) else: log.error("Labels config verification failed.") metrics.incr("error_count", tags=["type:misconfig", f"config_type:{config_type}"]) sys.exit(1) # Default logic removed for now because it will create hundreds of budgets # And we need to decide good thresholds for this ''' default_dict = budget_dict['default'] config_type = PLUTUS_CONFIG_TYPE_DEFAULT if verify_default_yaml(default_dict): # Query for all projects. If no budget exists for project, add a default budget for p in resource_manager_client.list_projects(): project_id = p.project_id project_number = gcp.get_project_number(project_id) if project_number is not None: budgets = gcp.get_budgets_by_project(project, project_number) if len(budgets) == 0: # No budget exists for this project, so we create one default_dict['project_id'] = project_id # TODO - add and test # project = ProjectBudget(default_dict, config_type, # billing_account_id, default_pubsub_topic) log.info(f"Creating default budget for plutus-default-{project_id}") # budget = gcp.create_budget(project) # metrics.incr("default_budget_created_count", tags=[]) #if budget is not None: # with mysql_conn.cursor() as mysql_cursor: # upsert_budget(mysql_cursor, budget, project.project_id, # config_type, project.alert_emails) # TODO - delete plutus-default budget if > 1 budget elif len(budgets) > 1: # TODO - if one of the budgets is a plutus-default, delete via API and in Mysql pass else: log.error("Default config verification failed.") metrics.incr("error_count", tags=["type:misconfig", f"config_type:{config_type}"]) sys.exit(1) ''' log.info("Plutus run complete.")