Add tests for billing sample (GoogleCloudPlatform#7432)

Author: busunkim96

billing/main.py

@@ -14,64 +14,59 @@
 
 import base64
 import json
-import os
-from googleapiclient import discovery
-PROJECT_ID = os.getenv('GCP_PROJECT')
-PROJECT_NAME = f'projects/{PROJECT_ID}'
-def stop_billing(data, context):
-    pubsub_data = base64.b64decode(data['data']).decode('utf-8')
-    pubsub_json = json.loads(pubsub_data)
-    cost_amount = pubsub_json['costAmount']
-    budget_amount = pubsub_json['budgetAmount']
-    if cost_amount <= budget_amount:
-        print(f'No action necessary. (Current cost: {cost_amount})')
-        return
 
-    if PROJECT_ID is None:
-        print('No project specified with environment variable')
-        return
+import google.auth
+from google.cloud import billing
 
-    billing = discovery.build(
-        'cloudbilling',
-        'v1',
-        cache_discovery=False,
-    )
 
-    projects = billing.projects()
+PROJECT_ID = google.auth.default()[1]
+cloud_billing_client = billing.CloudBillingClient()
+
+
+def stop_billing(data: dict, context):
+    pubsub_data = base64.b64decode(data["data"]).decode("utf-8")
+    pubsub_json = json.loads(pubsub_data)
+    cost_amount = pubsub_json["costAmount"]
+    budget_amount = pubsub_json["budgetAmount"]
+    if cost_amount <= budget_amount:
+        print(f"No action necessary. (Current cost: {cost_amount})")
+        return
 
-    billing_enabled = __is_billing_enabled(PROJECT_NAME, projects)
+    project_name = cloud_billing_client.common_project_path(PROJECT_ID)
+    billing_enabled = _is_billing_enabled(project_name)
 
     if billing_enabled:
-        __disable_billing_for_project(PROJECT_NAME, projects)
+        _disable_billing_for_project(project_name)
     else:
-        print('Billing already disabled')
+        print("Billing already disabled")
 
 
-def __is_billing_enabled(project_name, projects):
-    """
-    Determine whether billing is enabled for a project
-    @param {string} project_name Name of project to check if billing is enabled
-    @return {bool} Whether project has billing enabled or not
-    """
-    try:
-        res = projects.getBillingInfo(name=project_name).execute()
-        return res['billingEnabled']
-    except KeyError:
-        # If billingEnabled isn't part of the return, billing is not enabled
-        return False
-    except Exception:
-        print('Unable to determine if billing is enabled on specified project, assuming billing is enabled')
-        return True
-
-
-def __disable_billing_for_project(project_name, projects):
+def _is_billing_enabled(project_name: str) -> bool:
+    """Determine whether billing is enabled for a project
+
+    Args:
+        project_name (str): Name of project to check if billing is enabled
+
+    Returns:
+        bool: Whether project has billing enabled or not
     """
-    Disable billing for a project by removing its billing account
-    @param {string} project_name Name of project disable billing on
+    request = billing.GetProjectBillingInfoRequest(name=project_name)
+    project_billing_info = cloud_billing_client.get_project_billing_info(request)
+
+    return project_billing_info.billing_enabled
+
+
+def _disable_billing_for_project(project_name: str) -> None:
+    """Disable billing for a project by removing its billing account
+
+    Args:
+        project_name (str): Name of project disable billing on
     """
-    body = {'billingAccountName': ''}  # Disable billing
-    try:
-        res = projects.updateBillingInfo(name=project_name, body=body).execute()
-        print(f'Billing disabled: {json.dumps(res)}')
-    except Exception:
-        print('Failed to disable billing, possibly check permissions')
+    request = billing.UpdateProjectBillingInfoRequest(
+        name=project_name,
+        project_billing_info=billing.ProjectBillingInfo(
+            billing_account_name=""  # Disable billing
+        ),
+    )
+    project_biling_info = cloud_billing_client.update_project_billing_info(request)
+    print(f"Billing disabled: {project_biling_info}")

billing/requirements-test.txt

@@ -0,0 +1 @@
+pytest==6.2.4

billing/requirements.txt

@@ -0,0 +1 @@
+google-cloud-billing==1.4.1

billing/test_main.py

@@ -0,0 +1,102 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import json
+from unittest import mock
+
+import google.auth
+from google.cloud import billing
+
+from main import _is_billing_enabled, stop_billing
+
+
+PROJECT_ID = google.auth.default()[1]
+
+# NOTE(busunkim): These tests use mocks instead of disabling/enabling
+# the test project because a service account cannot be
+# granted sufficient permissions to add a biling account to the project.
+# https://cloud.google.com/billing/docs/how-to/modify-project#enable_billing_for_an_existing_project
+
+
+def test__is_billing_enabled():
+    assert _is_billing_enabled(f"projects/{PROJECT_ID}")
+
+
+def test_stop_billing_under_budget(capsys):
+    billing_data = {
+        "costAmount": 10,
+        "budgetAmount": 100.1,
+    }
+
+    encoded_data = base64.b64encode(json.dumps(billing_data).encode("utf-8")).decode(
+        "utf-8"
+    )
+
+    pubsub_message = {"data": encoded_data}
+    stop_billing(pubsub_message, None)
+    stdout, _ = capsys.readouterr()
+
+    assert "No action necessary" in stdout
+
+
+def test_stop_billing_over_budget(capsys):
+    billing_data = {
+        "costAmount": 120,
+        "budgetAmount": 100.1,
+    }
+
+    encoded_data = base64.b64encode(json.dumps(billing_data).encode("utf-8")).decode(
+        "utf-8"
+    )
+
+    pubsub_message = {"data": encoded_data}
+
+    with mock.patch(
+        "google.cloud.billing.CloudBillingClient.update_project_billing_info",
+        autospec=True,
+    ) as update_billing:
+        stop_billing(pubsub_message, None)
+        update_billing.assert_called_once()
+
+        assert update_billing.call_args[0][1].name == f"projects/{PROJECT_ID}"
+        assert (
+            update_billing.call_args[0][1].project_billing_info.billing_account_name
+            == ""
+        )
+    stdout, _ = capsys.readouterr()
+    assert "Billing disabled" in stdout
+
+
+def test_stop_billing_already_disabled(capsys):
+    billing_data = {
+        "costAmount": 120,
+        "budgetAmount": 100.1,
+    }
+
+    encoded_data = base64.b64encode(json.dumps(billing_data).encode("utf-8")).decode(
+        "utf-8"
+    )
+
+    pubsub_message = {"data": encoded_data}
+
+    with mock.patch(
+        "google.cloud.billing.CloudBillingClient.get_project_billing_info",
+        autospec=True,
+        return_value=billing.ProjectBillingInfo(billing_enabled=False),
+    ):
+        stop_billing(pubsub_message, None)
+
+    stdout, _ = capsys.readouterr()
+    assert "Billing already disabled" in stdout