chore(samples): fix LRO wrapper (#99)

packages/google-cloud-compute/samples/snippets/quickstart.py

@@ -121,6 +121,7 @@ def create_instance(
         Instance object.
     """
     instance_client = compute_v1.InstancesClient()
+    operation_client = compute_v1.ZoneOperationsClient()
 
     # Describe the size and source image of the boot disk to attach to the instance.
     disk = compute_v1.AttachedDisk()
@@ -155,8 +156,7 @@ def create_instance(
     # Wait for the create operation to complete.
     print(f"Creating the {instance_name} instance in {zone}...")
     operation = instance_client.insert(request=request)
-    if operation.status == compute_v1.Operation.Status.RUNNING:
-        operation_client = compute_v1.ZoneOperationsClient()
+    while operation.status != compute_v1.Operation.Status.DONE:
         operation = operation_client.wait(
             operation=operation.name, zone=zone, project=project_id
         )
@@ -180,13 +180,13 @@ def delete_instance(project_id: str, zone: str, machine_name: str) -> None:
         machine_name: name of the machine you want to delete.
     """
     instance_client = compute_v1.InstancesClient()
+    operation_client = compute_v1.ZoneOperationsClient()
 
     print(f"Deleting {machine_name} from {zone}...")
     operation = instance_client.delete(
         project=project_id, zone=zone, instance=machine_name
     )
-    if operation.status == compute_v1.Operation.Status.RUNNING:
-        operation_client = compute_v1.ZoneOperationsClient()
+    while operation.status != compute_v1.Operation.Status.DONE:
         operation = operation_client.wait(
             operation=operation.name, zone=zone, project=project_id
         )

packages/google-cloud-compute/samples/snippets/sample_default_values.py

@@ -59,7 +59,12 @@ def set_usage_export_bucket(project_id: str, bucket_name: str,
         project=project_id, usage_export_location_resource=usage_export_location)
 
     op_client = compute_v1.GlobalOperationsClient()
-    op_client.wait(project=project_id, operation=operation.name)
+
+    while operation.status != compute_v1.Operation.Status.DONE:
+        operation = op_client.wait(
+            operation=operation.name, project=project_id
+        )
+
 # [END compute_usage_report_set]
 
 
@@ -113,5 +118,9 @@ def disable_usage_export(project_id: str) -> None:
         project=project_id, usage_export_location_resource=None)
 
     op_client = compute_v1.GlobalOperationsClient()
-    op_client.wait(project=project_id, operation=operation.name)
+
+    while operation.status != compute_v1.Operation.Status.DONE:
+        operation = op_client.wait(
+            operation=operation.name, project=project_id
+        )
 # [END compute_usage_report_disable]

packages/google-cloud-compute/samples/snippets/sample_start_stop.py

@@ -35,7 +35,10 @@ def start_instance(project_id: str, zone: str, instance_name: str):
 
     op = instance_client.start(project=project_id, zone=zone, instance=instance_name)
 
-    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    while op.status != compute_v1.Operation.Status.DONE:
+        op = op_client.wait(
+            operation=op.name, zone=zone, project=project_id
+        )
     return
 # [END compute_start_instance]
 
@@ -71,7 +74,10 @@ def start_instance_with_encryption_key(project_id: str, zone: str, instance_name
     op = instance_client.start_with_encryption_key(project=project_id, zone=zone, instance=instance_name,
                                                    instances_start_with_encryption_key_request_resource=enc_data)
 
-    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    while op.status != compute_v1.Operation.Status.DONE:
+        op = op_client.wait(
+            operation=op.name, zone=zone, project=project_id
+        )
     return
 # [END compute_start_enc_instance]
 
@@ -91,7 +97,10 @@ def stop_instance(project_id: str, zone: str, instance_name: str):
 
     op = instance_client.stop(project=project_id, zone=zone, instance=instance_name)
 
-    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    while op.status != compute_v1.Operation.Status.DONE:
+        op = op_client.wait(
+            operation=op.name, zone=zone, project=project_id
+        )
     return
 # [END compute_stop_instance]
 
@@ -111,6 +120,9 @@ def reset_instance(project_id: str, zone: str, instance_name: str):
 
     op = instance_client.reset(project=project_id, zone=zone, instance=instance_name)
 
-    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    while op.status != compute_v1.Operation.Status.DONE:
+        op = op_client.wait(
+            operation=op.name, zone=zone, project=project_id
+        )
     return
 # [END compute_reset_instance]

packages/google-cloud-compute/samples/snippets/test_sample_start_stop.py

@@ -77,7 +77,10 @@ def _create_instance(request: compute_v1.InsertInstanceRequest):
     operation_client = compute_v1.ZoneOperationsClient()
 
     operation = instance_client.insert(request=request)
-    operation_client.wait(operation=operation.name, zone=INSTANCE_ZONE, project=PROJECT)
+    while operation.status != compute_v1.Operation.Status.DONE:
+        operation = operation_client.wait(
+            operation=operation.name, zone=INSTANCE_ZONE, project=PROJECT
+        )
 
     return instance_client.get(project=PROJECT, zone=INSTANCE_ZONE, instance=request.instance_resource.name)
 
@@ -87,7 +90,10 @@ def _delete_instance(instance: compute_v1.Instance):
     operation_client = compute_v1.ZoneOperationsClient()
 
     operation = instance_client.delete(project=PROJECT, zone=INSTANCE_ZONE, instance=instance.name)
-    operation_client.wait(operation=operation.name, zone=INSTANCE_ZONE, project=PROJECT)
+    while operation.status != compute_v1.Operation.Status.DONE:
+        operation = operation_client.wait(
+            operation=operation.name, zone=INSTANCE_ZONE, project=PROJECT
+        )
 
 
 def _get_status(instance: compute_v1.Instance) -> compute_v1.Instance.Status: