Cloud client library migration for registry management features (GoogleCloudPlatform#2453)

gguuss · web-flow · commit 223abdede962 · 2019-10-09T17:29:14.000-07:00
* Cloud client library migration for registry management features
diff --git a/iot/api-client/manager/manager.py b/iot/api-client/manager/manager.py
@@ -35,6 +35,7 @@
 import sys
 import time
 
+from google.api_core.exceptions import AlreadyExists
 from google.cloud import iot_v1
 from google.cloud import pubsub
 from google.oauth2 import service_account
@@ -143,37 +144,29 @@ def create_device(
     """Create a device to bind to a gateway if it does not exist."""
     # [START iot_create_device]
     # Check that the device doesn't already exist
+    client = iot_v1.DeviceManagerClient()
+
     exists = False
 
-    client = get_client(service_account_json)
-    registry_path = 'projects/{}/locations/{}/registries/{}'.format(
-        project_id, cloud_region, registry_id)
+    parent = client.registry_path(project_id, cloud_region, registry_id)
 
-    devices = client.projects().locations().registries().devices(
-        ).list(
-                parent=registry_path, fieldMask='config,gatewayConfig'
-        ).execute().get('devices', [])
+    devices = list(client.list_devices(parent=parent))
 
     for device in devices:
-        if device.get('id') == device_id:
+        if device.id == device_id:
             exists = True
 
     # Create the device
-    registry_name = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
-
     device_template = {
         'id': device_id,
-        'gatewayConfig': {
-          'gatewayType': 'NON_GATEWAY',
-          'gatewayAuthMethod': 'ASSOCIATION_ONLY'
+        'gateway_config': {
+          'gateway_type': 'NON_GATEWAY',
+          'gateway_auth_method': 'ASSOCIATION_ONLY'
         }
     }
-    devices = client.projects().locations().registries().devices()
 
     if not exists:
-        res = devices.create(
-                parent=registry_name, body=device_template).execute()
+        res = client.create_device(parent, device_template)
         print('Created Device {}'.format(res))
     else:
         print('Device exists, skipping')
@@ -217,12 +210,17 @@ def delete_registry(
     """Deletes the specified registry."""
     # [START iot_delete_registry]
     print('Delete registry')
-    client = get_client(service_account_json)
-    registry_name = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
 
-    registries = client.projects().locations().registries()
-    return registries.delete(name=registry_name).execute()
+    client = iot_v1.DeviceManagerClient()
+    registry_path = client.registry_path(project_id, cloud_region, registry_id)
+
+    try:
+        response = client.delete_device_registry(registry_path)
+        print('Deleted registry')
+        return response
+    except HttpError:
+        print('Error, registry not deleted')
+        return ""
     # [END iot_delete_registry]
 
 
@@ -330,40 +328,40 @@ def create_registry(
     """ Creates a registry and returns the result. Returns an empty result if
     the registry already exists."""
     # [START iot_create_registry]
-    client = get_client(service_account_json)
-    registry_parent = 'projects/{}/locations/{}'.format(
-            project_id,
-            cloud_region)
+    client = iot_v1.DeviceManagerClient()
+    parent = client.location_path(project_id, cloud_region)
+
+    if not pubsub_topic.startswith('projects/'):
+        pubsub_topic = 'projects/{}/topics/{}'.format(project_id, pubsub_topic)
+
     body = {
-        'eventNotificationConfigs': [{
-            'pubsubTopicName': pubsub_topic
+        'event_notification_configs': [{
+            'pubsub_topic_name': pubsub_topic
         }],
         'id': registry_id
     }
-    request = client.projects().locations().registries().create(
-        parent=registry_parent, body=body)
 
     try:
-        response = request.execute()
+        response = client.create_device_registry(parent, body)
         print('Created registry')
         return response
     except HttpError:
         print('Error, registry not created')
         return ""
+    except AlreadyExists:
+        print('Error, registry already exists')
+        return ""
     # [END iot_create_registry]
 
 
 def get_registry(
         service_account_json, project_id, cloud_region, registry_id):
     """ Retrieves a device registry."""
     # [START iot_get_registry]
-    client = get_client(service_account_json)
-    registry_parent = 'projects/{}/locations/{}'.format(
-            project_id,
-            cloud_region)
-    topic_name = '{}/registries/{}'.format(registry_parent, registry_id)
-    request = client.projects().locations().registries().get(name=topic_name)
-    return request.execute()
+    client = iot_v1.DeviceManagerClient()
+    registry_path = client.registry_path(project_id, cloud_region, registry_id)
+
+    return client.get_device_registry(registry_path)
     # [END iot_get_registry]
 
 
@@ -377,7 +375,7 @@ def open_registry(
         service_account_json, project_id, cloud_region,
         pubsub_topic, registry_id)
 
-    if response == "":
+    if response == '':
         # Device registry already exists
         print(
             'Registry {} already exists - looking it up instead.'.format(
@@ -386,7 +384,7 @@ def open_registry(
             service_account_json, project_id, cloud_region,
             registry_id)
 
-    print('Registry {} opened: '.format(response.get('name')))
+    print('Registry {} opened: '.format(response.name))
     print(response)
 
 
@@ -512,12 +510,11 @@ def get_iam_permissions(
         service_account_json, project_id, cloud_region, registry_id):
     """Retrieves IAM permissions for the given registry."""
     # [START iot_get_iam_policy]
-    client = get_client(service_account_json)
-    registry_path = 'projects/{}/locations/{}/registries/{}'.format(
-        project_id, cloud_region, registry_id)
+    client = iot_v1.DeviceManagerClient()
+
+    registry_path = client.registry_path(project_id, cloud_region, registry_id)
 
-    policy = client.projects().locations().registries().getIamPolicy(
-            resource=registry_path, body={}).execute()
+    policy = client.get_iam_policy(registry_path)
 
     return policy
     # [END iot_get_iam_policy]
@@ -528,23 +525,18 @@ def set_iam_permissions(
         member):
     """Sets IAM permissions for the given registry to a single role/member."""
     # [START iot_set_iam_policy]
-    client = get_client(service_account_json)
+    client = iot_v1.DeviceManagerClient()
+    registry_path = client.registry_path(project_id, cloud_region, registry_id)
 
-    registry_path = 'projects/{}/locations/{}/registries/{}'.format(
-        project_id, cloud_region, registry_id)
     body = {
-        "policy":
-        {
-            "bindings":
-            [{
-                "members": [member],
-                "role": role
-            }]
-        }
+        'bindings':
+        [{
+            'members': [member],
+            'role': role
+        }]
     }
 
-    return client.projects().locations().registries().setIamPolicy(
-            resource=registry_path, body=body).execute()
+    return client.set_iam_policy(registry_path, body)
     # [END iot_set_iam_policy]
 
 
@@ -577,17 +569,13 @@ def create_gateway(
     # [START iot_create_gateway]
     # Check that the gateway doesn't already exist
     exists = False
-    client = get_client(service_account_json)
-    registry_path = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
+    client = iot_v1.DeviceManagerClient()
 
-    devices = client.projects().locations().registries().devices(
-            ).list(
-                    parent=registry_path, fieldMask='config,gatewayConfig'
-            ).execute().get('devices', [])
+    parent = client.registry_path(project_id, cloud_region, registry_id)
+    devices = list(client.list_devices(parent=parent))
 
     for device in devices:
-        if device.get('id') == gateway_id:
+        if device.id == gateway_id:
             exists = True
         print('Device: {} : {} : {} : {}'.format(
             device.get('id'),
@@ -596,10 +584,6 @@ def create_gateway(
             device.get('gatewayConfig')
             ))
 
-    # Create the gateway
-    registry_name = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
-
     with io.open(certificate_file) as f:
         certificate = f.read()
 
@@ -617,17 +601,15 @@ def create_gateway(
                 'key': certificate
             }
         }],
-        'gatewayConfig': {
-          'gatewayType': 'GATEWAY',
-          'gatewayAuthMethod': 'ASSOCIATION_ONLY'
+        'gateway_config': {
+          'gateway_type': 'GATEWAY',
+          'gateway_auth_method': 'ASSOCIATION_ONLY'
         }
     }
-    devices = client.projects().locations().registries().devices()
 
     if not exists:
-        res = devices.create(
-                parent=registry_name, body=device_template).execute()
-        print('Created gateway {}'.format(res))
+        res = client.create_device(parent, device_template)
+        print('Created Gateway {}'.format(res))
     else:
         print('Gateway exists, skipping')
     # [END iot_create_gateway]
@@ -638,21 +620,17 @@ def bind_device_to_gateway(
         gateway_id):
     """Binds a device to a gateway."""
     # [START iot_bind_device_to_gateway]
-    client = get_client(service_account_json)
+    client = iot_v1.DeviceManagerClient()
 
     create_device(
             service_account_json, project_id, cloud_region, registry_id,
             device_id)
 
-    registry_name = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
-    bind_request = {
-        'deviceId': device_id,
-        'gatewayId': gateway_id
-    }
-    client.projects().locations().registries().bindDeviceToGateway(
-            parent=registry_name, body=bind_request).execute()
-    print('Device Bound!')
+    parent = client.registry_path(project_id, cloud_region, registry_id)
+
+    res = client.bind_device_to_gateway(parent, gateway_id, device_id)
+
+    print('Device Bound! {}'.format(res))
     # [END iot_bind_device_to_gateway]
 
 
@@ -661,17 +639,12 @@ def unbind_device_from_gateway(
         gateway_id):
     """Unbinds a device to a gateway."""
     # [START iot_unbind_device_from_gateway]
-    client = get_client(service_account_json)
+    client = iot_v1.DeviceManagerClient()
 
-    registry_name = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
-    bind_request = {
-        'deviceId': device_id,
-        'gatewayId': gateway_id
-    }
+    parent = client.registry_path(project_id, cloud_region, registry_id)
+
+    res = client.unbind_device_from_gateway(parent, gateway_id, device_id)
 
-    res = client.projects().locations().registries().unbindDeviceFromGateway(
-        parent=registry_name, body=bind_request).execute()
     print('Device unbound: {}'.format(res))
     # [END iot_unbind_device_from_gateway]
 
@@ -680,19 +653,18 @@ def list_gateways(
         service_account_json, project_id, cloud_region, registry_id):
     """Lists gateways in a registry"""
     # [START iot_list_gateways]
-    client = get_client(service_account_json)
-    registry_path = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
+    client = iot_v1.DeviceManagerClient()
 
-    devices = client.projects().locations().registries().devices(
-            ).list(
-                    parent=registry_path, fieldMask='config,gatewayConfig'
-            ).execute().get('devices', [])
+    path = client.registry_path(project_id, cloud_region, registry_id)
+    mask = iot_v1.types.FieldMask()
+    mask.paths.append('config')
+    mask.paths.append('gateway_config')
+    devices = list(client.list_devices(parent=path, field_mask=mask))
 
     for device in devices:
-        if device.get('gatewayConfig') is not None:
-            if device.get('gatewayConfig').get('gatewayType') == 'GATEWAY':
-                print('Gateway ID: {}\n\t{}'.format(device.get('id'), device))
+        if device.gateway_config is not None:
+            if device.gateway_config.gateway_type == 1:
+                print('Gateway ID: {}\n\t{}'.format(device.id, device))
     # [END iot_list_gateways]
 
 
@@ -701,23 +673,18 @@ def list_devices_for_gateway(
         gateway_id):
     """List devices bound to a gateway"""
     # [START iot_list_devices_for_gateway]
-    client = get_client(service_account_json)
+    client = iot_v1.DeviceManagerClient()
 
-    registry_name = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
+    path = client.registry_path(project_id, cloud_region, registry_id)
 
-    devices = client.projects().locations().registries().devices(
-        ).list(
-                parent=registry_name,
-                gatewayListOptions_associationsGatewayId=gateway_id
-        ).execute()
+    devices = list(client.list_devices(
+        parent=path,
+        gateway_list_options={'associations_gateway_id': gateway_id}))
 
     found = False
-    for device in devices.get('devices', []):
+    for device in devices:
         found = True
-        print('Device: {} : {}'.format(
-                device.get('numId'),
-                device.get('id')))
+        print('Device: {} : {}'.format(device.num_id, device.id))
 
     if not found:
         print('No devices bound to gateway {}'.format(gateway_id))
diff --git a/iot/api-client/manager/manager_test.py b/iot/api-client/manager/manager_test.py
@@ -137,7 +137,7 @@ def test_create_delete_registry(test_topic, capsys):
 
     # Check that create / list worked
     assert 'Created registry' in out
-    assert 'eventNotificationConfig' in out
+    assert 'event_notification_config' in out
 
     # Clean up
     manager.delete_registry(
@@ -167,7 +167,8 @@ def test_get_iam_permissions(test_topic, capsys):
 
     # Check that create / list worked
     assert 'Created registry' in out
-    assert 'eventNotificationConfig' in out
+    assert 'event_notification_config' in out
+    assert 'dpebot' in out
     assert 'etag' in out
 
     # Clean up
@@ -430,7 +431,7 @@ def test_create_gateway(test_topic, capsys):
 
     out, _ = capsys.readouterr()
 
-    assert 'Created gateway' in out
+    assert 'Created Gateway' in out
 
 
 def test_list_gateways(test_topic, capsys):
