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