Update compute metadata hanging get example to use "wait_for_change".

compute/metadata/main.py

@@ -26,44 +26,52 @@
 import requests
 
 
-METADATA_URL = "http://metadata.google.internal/computeMetadata/v1/"
+METADATA_URL = 'http://metadata.google.internal/computeMetadata/v1/'
 METADATA_HEADERS = {'Metadata-Flavor': 'Google'}
 
 
 def wait_for_maintenance(callback):
  url = METADATA_URL + 'instance/maintenance-event'
-last_in_maintenance = False
+last_maintenance_event = None
  # [START hanging_get]
- last_etag = 0
+ last_etag = "0"
 
  while True:
  r = requests.get(
  url,
- params={'last_etag': last_etag},
+ params={'last_etag': last_etag, 'wait_for_change': True},
  headers=METADATA_HEADERS)
 
  # During maintenance the service can return a 503, so these should
  # be retried.
  if r.status_code == 503:
  time.sleep(1)
  continue
+ elif r.status_code != requests.codes.ok:
+ print('Unexpected HTTP return code from metadata server: {}:{}.'
+ .format(r.status_code, r.reason))
+ time.sleep(1)
+ continue
 
  last_etag = r.headers['etag']
  # [END hanging_get]
 
- if r.text == 'MIGRATE_ON_HOST_MAINTENANCE':
- in_maintenance = True
+ if r.text != 'NONE':
+ # Possible events:
+ # MIGRATE_ON_HOST_MAINTENANCE: instance will be migrated
+ # SHUTDOWN_ON_HOST_MAINTENANCE: instance will be shut down
+ maintenance_event = r.text
  else:
-in_maintenance = False
+maintenance_event = None
 
- if in_maintenance != last_in_maintenance:
-last_in_maintenance = in_maintenance
- callback(in_maintenance)
+ if maintenance_event != last_maintenance_event:
+last_maintenance_event = maintenance_event
+ callback(maintenance_event)
 
 
-def maintenance_callback(status):
- if status:
- print('Undergoing host maintenance')
+def maintenance_callback(event):
+ if event:
+ print('Undergoing host maintenance: {}'.format(event))
  else:
  print('Finished host maintenance')
 

compute/metadata/main_test.py

@@ -13,6 +13,7 @@
 
 import main
 import mock
+import requests
 
 
 @mock.patch('main.requests')
@@ -31,6 +32,7 @@ def test_wait_for_maintenance(requests_mock):
  response3_mock = mock.Mock()
  response3_mock.status_code = 503
 
+ requests_mock.codes.ok = requests.codes.ok
  requests_mock.get.side_effect = [
  response1_mock, response2_mock, response3_mock, response2_mock,
  StopIteration()]
@@ -43,5 +45,6 @@ def test_wait_for_maintenance(requests_mock):
  pass
 
  assert callback_mock.call_count == 2
- assert callback_mock.call_args_list[0][0] == (True,)
- assert callback_mock.call_args_list[1][0] == (False,)
+ assert callback_mock.call_args_list[0][0] == (
+ 'MIGRATE_ON_HOST_MAINTENANCE',)
+ assert callback_mock.call_args_list[1][0] == (None,)