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

…#333)

We need to include "wait_for_change" in order to perform a hanging get. In
addition, we should deal with possible errors (like too many hanging gets).
Finally, this change allows for any of the maintenance event supported by
the metadata server.

compute/metadata/main.py

@@ -26,44 +26,48 @@
 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
+        r.raise_for_status()
 
         last_etag = r.headers['etag']
         # [END hanging_get]
 
-        if r.text == 'MIGRATE_ON_HOST_MAINTENANCE':
-            in_maintenance = True
+        if r.text == 'NONE':
+            maintenance_event = None
         else:
-            in_maintenance = False
+            # Possible events:
+            #   MIGRATE_ON_HOST_MAINTENANCE: instance will be migrated
+            #   SHUTDOWN_ON_HOST_MAINTENANCE: instance will be shut down
+            maintenance_event = r.text
 
-        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,)