Add GAE billing tests + fix region tags (GoogleCloudPlatform#2774)

Ace Nassri · web-flow · commit 96a08deda94a · 2020-02-05T15:17:40.000-08:00
* Add GAE billing limit sample

* Fix lint (since I cant run it locally)

* Fix lint v2

* Add missing region tag blocks

* Add region tags to imports

* Add tests

* Actually save lint fixes :)

* Fix lint

* Add tests

* Actually save lint fixes :)

* Fix lint

* Only test GCF in Python 3.7

* GCF Py3.7 only, take 2

* DBG commit 1

* Remove DBG commit + capture Python version correctly
diff --git a/.kokoro/tests/run_tests.sh b/.kokoro/tests/run_tests.sh
@@ -76,6 +76,16 @@ for file in **/requirements.txt; do
         fi
     fi
 
+    # Skip unsupported Python versions for Cloud Functions
+    # (Some GCF samples' dependencies don't support them)
+    if [[ "$file" == "functions/"* ]]; then
+      PYTHON_VERSION="$(python --version 2>&1)"
+      if [[ "$PYTHON_VERSION" == "Python 2."* || "$PYTHON_VERSION" == "Python 3.5"* ]]; then
+        # echo -e "\n Skipping $file: Python $PYTHON_VERSION is not supported by Cloud Functions.\n"
+        continue
+      fi
+    fi
+
     echo "------------------------------------------------------------"
     echo "- testing $file"
     echo "------------------------------------------------------------"
diff --git a/functions/billing/main.py b/functions/billing/main.py
@@ -13,30 +13,36 @@
 # limitations under the License.
 
 # [START functions_billing_limit]
+# [START functions_billing_limit_appengine]
 # [START functions_billing_stop]
 import base64
 import json
 import os
 # [END functions_billing_stop]
 # [END functions_billing_limit]
+# [END functions_billing_limit_appengine]
 
 # [START functions_billing_limit]
+# [START functions_billing_limit_appengine]
 # [START functions_billing_stop]
 from googleapiclient import discovery
 from oauth2client.client import GoogleCredentials
 # [END functions_billing_stop]
 # [END functions_billing_limit]
+# [END functions_billing_limit_appengine]
 
 # [START functions_billing_slack]
 import slack
 # [END functions_billing_slack]
 
 # [START functions_billing_limit]
+# [START functions_billing_limit_appengine]
 # [START functions_billing_stop]
 PROJECT_ID = os.getenv('GCP_PROJECT')
 PROJECT_NAME = f'projects/{PROJECT_ID}'
 # [END functions_billing_stop]
 # [END functions_billing_limit]
+# [END functions_billing_limit_appengine]
 
 # [START functions_billing_slack]
 
diff --git a/functions/billing/main_test.py b/functions/billing/main_test.py
@@ -114,3 +114,50 @@ def discovery_mocker(x, *args, **kwargs):
     assert instances_mock.list.calledWith(project=PROJECT_ID, zone=ZONE)
     assert instances_mock.stop.call_count == 1
     assert instances_mock.execute.call_count == 2
+
+
+@patch('main.PROJECT_ID')
+@patch('main.ZONE')
+@patch('main.discovery')
+def test_limit_use_appengine(discovery_mock, ZONE, PROJECT_ID):
+    PROJECT_ID = 'my-project'
+    PROJECT_NAME = f'projects/{PROJECT_ID}'
+
+    data = {"budgetAmount": 400, "costAmount": 500}
+
+    pubsub_message = {
+        "data": base64.b64encode(bytes(json.dumps(data), 'utf-8')),
+        "attributes": {}
+    }
+
+    projects_mock = MagicMock()
+    projects_mock.projects = MagicMock(return_value=projects_mock)
+    projects_mock.getBillingInfo = MagicMock(return_value=projects_mock)
+    projects_mock.updateBillingInfo = MagicMock(return_value=projects_mock)
+
+    apps_list = [{'servingStatus': 'SERVING'}]
+    app_patch_mock = MagicMock()
+    apps_mock = MagicMock()
+    apps_mock.get.return_value.execute.return_value = apps_list
+    apps_mock.patch.return_value.execute = app_patch_mock
+    appengine_mock = MagicMock()
+    appengine_mock.apps.return_value = apps_mock
+
+    def discovery_mocker(x, *args, **kwargs):
+        if x == 'appengine':
+            return apps_mock
+        else:
+            return projects_mock
+
+    discovery_mock.build = MagicMock(side_effect=discovery_mocker)
+
+    main.limit_use_appengine(pubsub_message, None)
+
+    patch_body = {
+        'servingStatus': 'USER_DISABLED'
+    }
+
+    assert projects_mock.getBillingInfo.called_with(name=PROJECT_NAME)
+    assert apps_mock.get.calledWith(appsId=PROJECT_ID)
+    assert apps_mock.stop.calledWith(
+        appsId=PROJECT_ID, updateMask='serving_status', body=patch_body)
