Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 7e07ff8

Browse files
author
Ace Nassri
authored
GCF: Add samples for billing (GoogleCloudPlatform#1725)
* Add billing GCF samples Change-Id: Ie43de503c9600f23b7bb168e62948a01f868457d * Fix spacing Change-Id: Iaa73d0548b5196a9e002b330165fd09c172928da * Add token comment Change-Id: I537817dda32abb03629e8ebc4c9e45d69efb59f7
1 parent 23dfa12 commit 7e07ff8

File tree

3 files changed

+288
-0
lines changed

3 files changed

+288
-0
lines changed

‎functions/billing/main.py

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
# Copyright 2018 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the 'License');
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an 'AS IS' BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START functions_billing_limit]
16+
# [START functions_billing_stop]
17+
import base64
18+
import json
19+
# [END functions_billing_stop]
20+
import os
21+
# [END functions_billing_limit]
22+
23+
# [START functions_billing_limit]
24+
# [START functions_billing_stop]
25+
from googleapiclient import discovery
26+
from oauth2client.client import GoogleCredentials
27+
28+
# [END functions_billing_stop]
29+
# [END functions_billing_limit]
30+
31+
# [START functions_billing_slack]
32+
from slackclient import SlackClient
33+
# [END functions_billing_slack]
34+
35+
# [START functions_billing_limit]
36+
# [START functions_billing_stop]
37+
PROJECT_ID = os.getenv('GCP_PROJECT')
38+
PROJECT_NAME = f'projects/{PROJECT_ID}'
39+
# [END functions_billing_stop]
40+
# [END functions_billing_limit]
41+
42+
# [START functions_billing_slack]
43+
44+
# See https://api.slack.com/docs/token-types#bot for more info
45+
BOT_ACCESS_TOKEN = 'xxxx-111111111111-abcdefghidklmnopq'
46+
47+
CHANNEL = 'general'
48+
49+
slack_client = SlackClient(BOT_ACCESS_TOKEN)
50+
51+
52+
def notify_slack(data, context):
53+
pubsub_message = data
54+
55+
notification_attrs = json.dumps(pubsub_message['attributes'])
56+
notification_data = base64.b64decode(data['data']).decode('utf-8')
57+
budget_notification_text = f'{notification_attrs}, {notification_data}'
58+
59+
res = slack_client.api_call(
60+
'chat.postMessage',
61+
channel=CHANNEL,
62+
text=budget_notification_text)
63+
print(res)
64+
# [END functions_billing_slack]
65+
66+
67+
# [START functions_billing_limit]
68+
def stop_billing(data, context):
69+
pubsub_data = base64.b64decode(data['data']).decode('utf-8')
70+
pubsub_json = json.loads(pubsub_data)
71+
cost_amount = pubsub_json['costAmount']
72+
budget_amount = pubsub_json['budgetAmount']
73+
if cost_amount <= budget_amount:
74+
print(f'No action necessary. (Current cost: {cost_amount})')
75+
return
76+
77+
billing = discovery.build(
78+
'cloudbilling',
79+
'v1',
80+
cache_discovery=False,
81+
credentials=GoogleCredentials.get_application_default()
82+
)
83+
84+
projects = billing.projects()
85+
86+
if __is_billing_enabled(PROJECT_NAME, projects):
87+
print(__disable_billing_for_project(PROJECT_NAME, projects))
88+
else:
89+
print('Billing already disabled')
90+
91+
92+
def __is_billing_enabled(project_name, projects):
93+
"""
94+
Determine whether billing is enabled for a project
95+
@param {string} project_name Name of project to check if billing is enabled
96+
@return {bool} Whether project has billing enabled or not
97+
"""
98+
res = projects.getBillingInfo(name=project_name).execute()
99+
return res['billingEnabled']
100+
101+
102+
def __disable_billing_for_project(project_name, projects):
103+
"""
104+
Disable billing for a project by removing its billing account
105+
@param {string} project_name Name of project disable billing on
106+
@return {string} Text containing response from disabling billing
107+
"""
108+
body = {'billingAccountName': ''} # Disable billing
109+
res = projects.updateBillingInfo(name=project_name, body=body).execute()
110+
print(f'Billing disabled: {json.dumps(res)}')
111+
# [END functions_billing_stop]
112+
113+
114+
# [START functions_billing_limit]
115+
ZONE = 'us-west1-b'
116+
117+
118+
def limit_use(data, context):
119+
pubsub_data = base64.b64decode(data['data']).decode('utf-8')
120+
pubsub_json = json.loads(pubsub_data)
121+
cost_amount = pubsub_json['costAmount']
122+
budget_amount = pubsub_json['budgetAmount']
123+
if cost_amount <= budget_amount:
124+
print(f'No action necessary. (Current cost: {cost_amount})')
125+
return
126+
127+
compute = discovery.build(
128+
'compute',
129+
'v1',
130+
cache_discovery=False,
131+
credentials=GoogleCredentials.get_application_default()
132+
)
133+
instances = compute.instances()
134+
135+
instance_names = __list_running_instances(PROJECT_ID, ZONE, instances)
136+
__stop_instances(PROJECT_ID, ZONE, instance_names, instances)
137+
138+
139+
def __list_running_instances(project_id, zone, instances):
140+
"""
141+
@param {string} project_id ID of project that contains instances to stop
142+
@param {string} zone Zone that contains instances to stop
143+
@return {Promise} Array of names of running instances
144+
"""
145+
res = instances.list(project=project_id, zone=zone).execute()
146+
147+
items = res['items']
148+
running_names = [i['name'] for i in items if i['status'] == 'RUNNING']
149+
return running_names
150+
151+
152+
def __stop_instances(project_id, zone, instance_names, instances):
153+
"""
154+
@param {string} project_id ID of project that contains instances to stop
155+
@param {string} zone Zone that contains instances to stop
156+
@param {Array} instance_names Names of instance to stop
157+
@return {Promise} Response from stopping instances
158+
"""
159+
if not len(instance_names):
160+
print('No running instances were found.')
161+
return
162+
163+
for name in instance_names:
164+
instances.stop(
165+
project=project_id,
166+
zone=zone,
167+
instance=name).execute()
168+
print(f'Instance stopped successfully: {name}')
169+
# [END functions_billing_limit]

‎functions/billing/main_test.py

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# Copyright 2018, Google, LLC.
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
14+
import base64
15+
import json
16+
17+
from mock import MagicMock, patch
18+
19+
import main
20+
21+
22+
@patch('main.slack_client')
23+
def test_notify_slack(slack_client):
24+
slack_client.api_call = MagicMock()
25+
26+
data = {"budgetAmount": 400, "costAmount": 500}
27+
attrs = {"foo": "bar"}
28+
29+
pubsub_message = {
30+
"data": base64.b64encode(bytes(json.dumps(data), 'utf-8')),
31+
"attributes": attrs
32+
}
33+
34+
main.notify_slack(pubsub_message, None)
35+
36+
assert slack_client.api_call.called
37+
38+
39+
@patch('main.PROJECT_ID')
40+
@patch('main.discovery')
41+
def test_disable_billing(discovery_mock, PROJECT_ID):
42+
PROJECT_ID = 'my-project'
43+
PROJECT_NAME = f'projects/{PROJECT_ID}'
44+
45+
data = {"budgetAmount": 400, "costAmount": 500}
46+
47+
pubsub_message = {
48+
"data": base64.b64encode(bytes(json.dumps(data), 'utf-8')),
49+
"attributes": {}
50+
}
51+
52+
projects_mock = MagicMock()
53+
projects_mock.projects = MagicMock(return_value=projects_mock)
54+
projects_mock.getBillingInfo = MagicMock(return_value=projects_mock)
55+
projects_mock.updateBillingInfo = MagicMock(return_value=projects_mock)
56+
projects_mock.execute = MagicMock(return_value={'billingEnabled': True})
57+
58+
discovery_mock.build = MagicMock(return_value=projects_mock)
59+
60+
main.stop_billing(pubsub_message, None)
61+
62+
assert projects_mock.getBillingInfo.called_with(name=PROJECT_NAME)
63+
assert projects_mock.updateBillingInfo.called_with(
64+
name=PROJECT_NAME,
65+
body={'billingAccountName': ''}
66+
)
67+
assert projects_mock.execute.call_count == 2
68+
69+
70+
@patch('main.PROJECT_ID')
71+
@patch('main.ZONE')
72+
@patch('main.discovery')
73+
def test_limit_use(discovery_mock, ZONE, PROJECT_ID):
74+
PROJECT_ID = 'my-project'
75+
PROJECT_NAME = f'projects/{PROJECT_ID}'
76+
ZONE = 'my-zone'
77+
78+
data = {"budgetAmount": 400, "costAmount": 500}
79+
80+
pubsub_message = {
81+
"data": base64.b64encode(bytes(json.dumps(data), 'utf-8')),
82+
"attributes": {}
83+
}
84+
85+
instances_list = {
86+
"items": [
87+
{"name": "instance-1", "status": "RUNNING"},
88+
{"name": "instance-2", "status": "TERMINATED"}
89+
]
90+
}
91+
92+
instances_mock = MagicMock()
93+
instances_mock.instances = MagicMock(return_value=instances_mock)
94+
instances_mock.list = MagicMock(return_value=instances_mock)
95+
instances_mock.stop = MagicMock(return_value=instances_mock)
96+
instances_mock.execute = MagicMock(return_value=instances_list)
97+
98+
projects_mock = MagicMock()
99+
projects_mock.projects = MagicMock(return_value=projects_mock)
100+
projects_mock.getBillingInfo = MagicMock(return_value=projects_mock)
101+
projects_mock.execute = MagicMock(return_value={'billingEnabled': True})
102+
103+
def discovery_mocker(x, *args, **kwargs):
104+
if x == 'compute':
105+
return instances_mock
106+
else:
107+
return projects_mock
108+
109+
discovery_mock.build = MagicMock(side_effect=discovery_mocker)
110+
111+
main.limit_use(pubsub_message, None)
112+
113+
assert projects_mock.getBillingInfo.called_with(name=PROJECT_NAME)
114+
assert instances_mock.list.calledWith(project=PROJECT_ID, zone=ZONE)
115+
assert instances_mock.stop.call_count == 1
116+
assert instances_mock.execute.call_count == 2

‎functions/billing/requirements.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
slackclient==1.3.0
2+
oauth2client==4.1.3
3+
google-api-python-client==1.7.4

0 commit comments

Comments
 (0)