@@ -114,3 +114,50 @@ def discovery_mocker(x, *args, **kwargs):
114
114
assert instances_mock .list .calledWith (project = PROJECT_ID , zone = ZONE )
115
115
assert instances_mock .stop .call_count == 1
116
116
assert instances_mock .execute .call_count == 2
117
+
118
+
119
+ @patch ('main.PROJECT_ID' )
120
+ @patch ('main.ZONE' )
121
+ @patch ('main.discovery' )
122
+ def test_limit_use_appengine (discovery_mock , ZONE , PROJECT_ID ):
123
+ PROJECT_ID = 'my-project'
124
+ PROJECT_NAME = f'projects/{ PROJECT_ID } '
125
+
126
+ data = {"budgetAmount" : 400 , "costAmount" : 500 }
127
+
128
+ pubsub_message = {
129
+ "data" : base64 .b64encode (bytes (json .dumps (data ), 'utf-8' )),
130
+ "attributes" : {}
131
+ }
132
+
133
+ projects_mock = MagicMock ()
134
+ projects_mock .projects = MagicMock (return_value = projects_mock )
135
+ projects_mock .getBillingInfo = MagicMock (return_value = projects_mock )
136
+ projects_mock .updateBillingInfo = MagicMock (return_value = projects_mock )
137
+
138
+ apps_list = [{'servingStatus' : 'SERVING' }]
139
+ app_patch_mock = MagicMock ()
140
+ apps_mock = MagicMock ()
141
+ apps_mock .get .return_value .execute .return_value = apps_list
142
+ apps_mock .patch .return_value .execute = app_patch_mock
143
+ appengine_mock = MagicMock ()
144
+ appengine_mock .apps .return_value = apps_mock
145
+
146
+ def discovery_mocker (x , * args , ** kwargs ):
147
+ if x == 'appengine' :
148
+ return apps_mock
149
+ else :
150
+ return projects_mock
151
+
152
+ discovery_mock .build = MagicMock (side_effect = discovery_mocker )
153
+
154
+ main .limit_use_appengine (pubsub_message , None )
155
+
156
+ patch_body = {
157
+ 'servingStatus' : 'USER_DISABLED'
158
+ }
159
+
160
+ assert projects_mock .getBillingInfo .called_with (name = PROJECT_NAME )
161
+ assert apps_mock .get .calledWith (appsId = PROJECT_ID )
162
+ assert apps_mock .stop .calledWith (
163
+ appsId = PROJECT_ID , updateMask = 'serving_status' , body = patch_body )
0 commit comments