Remove stale test registries to avoid reaching limits (#2159)

gguuss · web-flow · commit 14b893cae651 · 2019-05-20T12:59:19.000-07:00
* Add cleanup step to tests to avoid hitting 100 registry limit
* Cleans up older test registries in case tests were cancelled and resources leaked
* Fixes broken style
* Renames includecode blocks to be ^iot(.*)
diff --git a/iot/api-client/manager/manager.py b/iot/api-client/manager/manager.py
@@ -142,22 +142,22 @@ def create_device(
         service_account_json, project_id, cloud_region, registry_id,
         device_id):
     """Create a device to bind to a gateway if it does not exist."""
-    # [START create_device]
+    # [START iot_create_device]
     # Check that the device doesn't already exist
     exists = False
 
     client = get_client(service_account_json)
     registry_path = 'projects/{}/locations/{}/registries/{}'.format(
-            project_id, cloud_region, registry_id)
+        project_id, cloud_region, registry_id)
 
     devices = client.projects().locations().registries().devices(
-            ).list(
-                    parent=registry_path, fieldMask='config,gatewayConfig'
-            ).execute().get('devices', [])
+        ).list(
+                parent=registry_path, fieldMask='config,gatewayConfig'
+        ).execute().get('devices', [])
 
     for device in devices:
-            if device.get('id') == device_id:
-                exists = True
+        if device.get('id') == device_id:
+            exists = True
 
     # Create the device
     registry_name = 'projects/{}/locations/{}/registries/{}'.format(
@@ -178,7 +178,7 @@ def create_device(
         print('Created Device {}'.format(res))
     else:
         print('Device exists, skipping')
-    # [END create_device]
+    # [END iot_create_device]
 
 
 def create_unauth_device(
@@ -295,9 +295,9 @@ def list_devices(
             ).list(parent=registry_path).execute().get('devices', [])
 
     for device in devices:
-            print('Device: {} : {}'.format(
-                    device.get('numId'),
-                    device.get('id')))
+        print('Device: {} : {}'.format(
+                device.get('numId'),
+                device.get('id')))
 
     return devices
     # [END iot_list_devices]
@@ -314,9 +314,9 @@ def list_registries(service_account_json, project_id, cloud_region):
         parent=registry_path).execute().get('deviceRegistries', [])
 
     for registry in registries:
-            print('id: {}\n\tname: {}'.format(
-                    registry.get('id'),
-                    registry.get('name')))
+        print('id: {}\n\tname: {}'.format(
+            registry.get('id'),
+            registry.get('name')))
 
     return registries
     # [END iot_list_registries]
@@ -375,7 +375,7 @@ def open_registry(
         service_account_json, project_id, cloud_region,
         pubsub_topic, registry_id)
 
-    if (response is ""):
+    if response == "":
         # Device registry already exists
         print(
             'Registry {} already exists - looking it up instead.'.format(
@@ -559,7 +559,7 @@ def create_gateway(
         service_account_json, project_id, cloud_region, registry_id, device_id,
         gateway_id, certificate_file, algorithm):
     """Create a gateway to bind devices to."""
-    # [START create_gateway]
+    # [START iot_create_gateway]
     # Check that the gateway doesn't already exist
     exists = False
     client = get_client(service_account_json)
@@ -572,14 +572,14 @@ def create_gateway(
             ).execute().get('devices', [])
 
     for device in devices:
-            if device.get('id') == gateway_id:
-                exists = True
-            print('Device: {} : {} : {} : {}'.format(
-                device.get('id'),
-                device.get('numId'),
-                device.get('config'),
-                device.get('gatewayConfig')
-                ))
+        if device.get('id') == gateway_id:
+            exists = True
+        print('Device: {} : {} : {} : {}'.format(
+            device.get('id'),
+            device.get('numId'),
+            device.get('config'),
+            device.get('gatewayConfig')
+            ))
 
     # Create the gateway
     registry_name = 'projects/{}/locations/{}/registries/{}'.format(
@@ -615,14 +615,14 @@ def create_gateway(
         print('Created gateway {}'.format(res))
     else:
         print('Gateway exists, skipping')
-    # [END create_gateway]
+    # [END iot_create_gateway]
 
 
 def bind_device_to_gateway(
         service_account_json, project_id, cloud_region, registry_id, device_id,
         gateway_id):
     """Binds a device to a gateway."""
-    # [START bind_device_to_gateway]
+    # [START iot_bind_device_to_gateway]
     client = get_client(service_account_json)
 
     create_device(
@@ -638,14 +638,14 @@ def bind_device_to_gateway(
     client.projects().locations().registries().bindDeviceToGateway(
             parent=registry_name, body=bind_request).execute()
     print('Device Bound!')
-    # [END bind_device_to_gateway]
+    # [END iot_bind_device_to_gateway]
 
 
 def unbind_device_from_gateway(
         service_account_json, project_id, cloud_region, registry_id, device_id,
         gateway_id):
     """Unbinds a device to a gateway."""
-    # [START unbind_device_from_gateway]
+    # [START iot_unbind_device_from_gateway]
     client = get_client(service_account_json)
 
     registry_name = 'projects/{}/locations/{}/registries/{}'.format(
@@ -658,13 +658,13 @@ def unbind_device_from_gateway(
     res = client.projects().locations().registries().unbindDeviceFromGateway(
         parent=registry_name, body=bind_request).execute()
     print('Device unbound: {}'.format(res))
-    # [END unbind_device_from_gateway]
+    # [END iot_unbind_device_from_gateway]
 
 
 def list_gateways(
         service_account_json, project_id, cloud_region, registry_id):
     """Lists gateways in a registry"""
-    # [START list_gateways]
+    # [START iot_list_gateways]
     client = get_client(service_account_json)
     registry_path = 'projects/{}/locations/{}/registries/{}'.format(
             project_id, cloud_region, registry_id)
@@ -678,14 +678,14 @@ def list_gateways(
         if device.get('gatewayConfig') is not None:
             if device.get('gatewayConfig').get('gatewayType') == 'GATEWAY':
                 print('Gateway ID: {}\n\t{}'.format(device.get('id'), device))
-    # [END list_gateways]
+    # [END iot_list_gateways]
 
 
 def list_devices_for_gateway(
         service_account_json, project_id, cloud_region, registry_id,
         gateway_id):
     """List devices bound to a gateway"""
-    # [START list_devices_for_gateway]
+    # [START iot_list_devices_for_gateway]
     client = get_client(service_account_json)
 
     registry_name = 'projects/{}/locations/{}/registries/{}'.format(
@@ -715,7 +715,7 @@ def list_devices_for_gateway(
     else:
         if not found:
             print('No devices bound to gateway {}'.format(gateway_id))
-    # [END list_devices_for_gateway]
+    # [END iot_list_devices_for_gateway]
 
 
 def parse_command_line_args():
diff --git a/iot/api-client/manager/manager_test.py b/iot/api-client/manager/manager_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import os
 import sys
 import time
@@ -40,6 +41,79 @@
 registry_id = 'test-registry-{}'.format(int(time.time()))
 
 
+@pytest.fixture(scope="session", autouse=True)
+def clean_up_registries():
+    all_registries = manager.list_registries(
+        service_account_json, project_id, cloud_region)
+
+    for registry in all_registries:
+        registry_id = registry.get('id')
+        if registry_id.find('test-registry-') == 0:
+            time_str = registry_id[
+                registry_id.rfind('-') + 1: len(registry_id)]
+            test_date = datetime.datetime.utcfromtimestamp(int(time_str))
+            now_date = datetime.datetime.utcfromtimestamp(int(time.time()))
+            difftime = now_date - test_date
+
+            # *NOTE* Restrict to registries used in the tests older than 30
+            #        days to prevent thrashing in the case of async tests
+            if (difftime.days > 30):
+                client = manager.get_client(service_account_json)
+                gateways = client.projects().locations().registries().devices(
+                    ).list(
+                        parent=registry.get('name'),
+                        fieldMask='config,gatewayConfig'
+                    ).execute().get('devices', [])
+                devices = client.projects().locations().registries().devices(
+                    ).list(parent=registry.get('name')).execute().get(
+                        'devices', [])
+
+                # Unbind devices from each gateway and delete
+                for gateway in gateways:
+                    gateway_id = gateway.get('id')
+                    bound = client.projects().locations().registries().devices(
+                        ).list(
+                            parent=registry.get('name'),
+                            gatewayListOptions_associationsGatewayId=gateway_id
+                        ).execute()
+                    if 'devices' in bound:
+                        for device in bound['devices']:
+                            bind_request = {
+                                'deviceId': device.get('id'),
+                                'gatewayId': gateway_id
+                            }
+                            client.projects().locations().registries(
+                                ).unbindDeviceFromGateway(
+                                    parent=registry.get('name'),
+                                    body=bind_request).execute()
+                    gateway_name = '{}/devices/{}'.format(
+                        registry.get('name'), gateway_id)
+                    client.projects().locations().registries().devices(
+                        ).delete(name=gateway_name).execute()
+
+                # Delete the devices
+                # Assumption is that the devices are not bound to gateways
+                for device in devices:
+                    device_name = '{}/devices/{}'.format(
+                        registry.get('name'), device.get('id'))
+                    print(device_name)
+                    remove_device = True
+                    try:
+                        client.projects().locations().registries().devices(
+                            ).get(name=device_name).execute()
+                    except Exception:
+                        remove_device = False
+
+                    if remove_device:
+                        print('removing {}'.format(device_name))
+                        client.projects().locations().registries().devices(
+                            ).delete(name=device_name).execute()
+
+                # Delete the old test registry
+                client.projects().locations().registries().delete(
+                    name=registry.get('name')).execute()
+
+
 @pytest.fixture(scope='module')
 def test_topic():
     topic = manager.create_iot_topic(project_id, topic_id)
diff --git a/iot/api-client/mqtt_example/cloudiot_mqtt_example.py b/iot/api-client/mqtt_example/cloudiot_mqtt_example.py
@@ -172,20 +172,20 @@ def get_client(
 
 def detach_device(client, device_id):
     """Detach the device from the gateway."""
-    # [START detach_device]
+    # [START iot_detach_device]
     detach_topic = '/devices/{}/detach'.format(device_id)
     print('Detaching: {}'.format(detach_topic))
     client.publish(detach_topic, '{}', qos=1)
-    # [END detach_device]
+    # [END iot_detach_device]
 
 
 def attach_device(client, device_id, auth):
     """Attach the device to the gateway."""
-    # [START attach_device]
+    # [START iot_attach_device]
     attach_topic = '/devices/{}/attach'.format(device_id)
     attach_payload = '{{"authorization" : "{}"}}'.format(auth)
     client.publish(attach_topic, attach_payload, qos=1)
-    # [END attach_device]
+    # [END iot_attach_device]
 
 
 def listen_for_messages(
@@ -194,7 +194,7 @@ def listen_for_messages(
         mqtt_bridge_hostname, mqtt_bridge_port, jwt_expires_minutes, duration,
         cb=None):
     """Listens for messages sent to the gateway and bound devices."""
-    # [START listen_for_messages]
+    # [START iot_listen_for_messages]
     global minimum_backoff_time
 
     jwt_iat = datetime.datetime.utcnow()
@@ -252,7 +252,7 @@ def listen_for_messages(
     detach_device(client, device_id)
 
     print('Finished.')
-    # [END listen_for_messages]
+    # [END iot_listen_for_messages]
 
 
 def send_data_from_bound_device(
