From 25406001bd45772978cc4d610e767a67c15c8c12 Mon Sep 17 00:00:00 2001 From: Jon Wayne Parrott Date: Tue, 18 Jul 2017 10:37:55 -0700 Subject: [PATCH] Update bigtable to use future.operation --- bigtable/google/cloud/bigtable/cluster.py | 37 +++++++++++----------- bigtable/google/cloud/bigtable/instance.py | 17 ++++------ bigtable/tests/system.py | 27 ++-------------- bigtable/tests/unit/test_cluster.py | 23 ++++++-------- bigtable/tests/unit/test_instance.py | 21 ++++++------ 5 files changed, 46 insertions(+), 79 deletions(-) diff --git a/bigtable/google/cloud/bigtable/cluster.py b/bigtable/google/cloud/bigtable/cluster.py index 8f7321b2a548..8d15547efae3 100644 --- a/bigtable/google/cloud/bigtable/cluster.py +++ b/bigtable/google/cloud/bigtable/cluster.py @@ -21,9 +21,7 @@ instance_pb2 as data_v2_pb2) from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) -from google.cloud.operation import Operation -from google.cloud.operation import register_type - +from google.cloud.future import operation _CLUSTER_NAME_RE = re.compile(r'^projects/(?P[^/]+)/' r'instances/(?P[^/]+)/clusters/' @@ -33,9 +31,6 @@ """Default number of nodes to use when creating a cluster.""" -register_type(messages_v2_pb2.UpdateClusterMetadata) - - def _prepare_create_request(cluster): """Creates a protobuf request for a CreateCluster request. @@ -208,15 +203,18 @@ def create(self): :returns: The long-running operation corresponding to the create operation. """ - request_pb = _prepare_create_request(self) - # We expect a `google.longrunning.operations_pb2.Operation`. client = self._instance._client + + # We expect a `google.longrunning.operations_pb2.Operation`. + request_pb = _prepare_create_request(self) operation_pb = client._instance_stub.CreateCluster(request_pb) - operation = Operation.from_pb(operation_pb, client) - operation.target = self - operation.caller_metadata['request_type'] = 'CreateCluster' - return operation + operation_future = operation.from_grpc( + operation_pb, + client._operations_stub, + data_v2_pb2.Cluster, + metadata_type=messages_v2_pb2.UpdateClusterMetadata) + return operation_future def update(self): """Update this cluster. @@ -236,18 +234,21 @@ def update(self): :returns: The long-running operation corresponding to the update operation. """ + client = self._instance._client + + # We expect a `google.longrunning.operations_pb2.Operation`. request_pb = data_v2_pb2.Cluster( name=self.name, serve_nodes=self.serve_nodes, ) - # We expect a `google.longrunning.operations_pb2.Operation`. - client = self._instance._client operation_pb = client._instance_stub.UpdateCluster(request_pb) - operation = Operation.from_pb(operation_pb, client) - operation.target = self - operation.caller_metadata['request_type'] = 'UpdateCluster' - return operation + operation_future = operation.from_grpc( + operation_pb, + client._operations_stub, + data_v2_pb2.Cluster, + metadata_type=messages_v2_pb2.UpdateClusterMetadata) + return operation_future def delete(self): """Delete this cluster. diff --git a/bigtable/google/cloud/bigtable/instance.py b/bigtable/google/cloud/bigtable/instance.py index 1de3cbcea814..958f16602953 100644 --- a/bigtable/google/cloud/bigtable/instance.py +++ b/bigtable/google/cloud/bigtable/instance.py @@ -26,8 +26,7 @@ from google.cloud.bigtable.cluster import Cluster from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES from google.cloud.bigtable.table import Table -from google.cloud.operation import Operation -from google.cloud.operation import register_type +from google.cloud.future import operation _EXISTING_INSTANCE_LOCATION_ID = 'see-existing-cluster' @@ -35,10 +34,6 @@ r'instances/(?P[a-z][-a-z0-9]*)$') -register_type(messages_v2_pb2.CreateInstanceMetadata) -register_type(data_v2_pb2.Instance) - - def _prepare_create_request(instance): """Creates a protobuf request for a CreateInstance request. @@ -232,10 +227,12 @@ def create(self): # We expect a `google.longrunning.operations_pb2.Operation`. operation_pb = self._client._instance_stub.CreateInstance(request_pb) - operation = Operation.from_pb(operation_pb, self._client) - operation.target = self - operation.caller_metadata['request_type'] = 'CreateInstance' - return operation + operation_future = operation.from_grpc( + operation_pb, + self._client._operations_stub, + data_v2_pb2.Instance, + metadata_type=messages_v2_pb2.CreateInstanceMetadata) + return operation_future def update(self): """Update this instance. diff --git a/bigtable/tests/system.py b/bigtable/tests/system.py index 1fcda808db39..cfc2cb17f805 100644 --- a/bigtable/tests/system.py +++ b/bigtable/tests/system.py @@ -32,7 +32,6 @@ from google.cloud.environment_vars import BIGTABLE_EMULATOR from test_utils.retry import RetryErrors -from test_utils.retry import RetryResult from test_utils.system import EmulatorCreds from test_utils.system import unique_resource_id @@ -65,27 +64,6 @@ class Config(object): IN_EMULATOR = False -def _wait_until_complete(operation, max_attempts=5): - """Wait until an operation has completed. - - :type operation: :class:`google.cloud.operation.Operation` - :param operation: Operation that has not completed. - - :type max_attempts: int - :param max_attempts: (Optional) The maximum number of times to check if - the operation has completed. Defaults to 5. - - :rtype: bool - :returns: Boolean indicating if the operation is complete. - """ - - def _operation_complete(result): - return result - - retry = RetryResult(_operation_complete, max_tries=max_attempts) - return retry(operation.poll)() - - def _retry_on_unavailable(exc): """Retry only errors whose status code is 'UNAVAILABLE'.""" from grpc import StatusCode @@ -117,8 +95,7 @@ def setUpModule(): # After listing, create the test instance. created_op = Config.INSTANCE.create() - if not _wait_until_complete(created_op): - raise RuntimeError('Instance creation exceed 5 seconds.') + created_op.result(timeout=10) def tearDownModule(): @@ -166,7 +143,7 @@ def test_create_instance(self): self.instances_to_delete.append(instance) # We want to make sure the operation completes. - self.assertTrue(_wait_until_complete(operation)) + operation.result(timeout=10) # Create a new instance instance and make sure it is the same. instance_alt = Config.CLIENT.instance(ALT_INSTANCE_ID, LOCATION_ID) diff --git a/bigtable/tests/unit/test_cluster.py b/bigtable/tests/unit/test_cluster.py index 3eb18f43863d..e244b55d6dff 100644 --- a/bigtable/tests/unit/test_cluster.py +++ b/bigtable/tests/unit/test_cluster.py @@ -15,6 +15,8 @@ import unittest +import mock + class TestCluster(unittest.TestCase): @@ -232,7 +234,7 @@ def test_reload(self): def test_create(self): from google.longrunning import operations_pb2 - from google.cloud.operation import Operation + from google.cloud.future import operation from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from tests.unit._testing import _FakeStub @@ -256,13 +258,9 @@ def test_create(self): # Perform the method and check the result. result = cluster.create() - self.assertIsInstance(result, Operation) - self.assertEqual(result.name, OP_NAME) - self.assertIs(result.target, cluster) - self.assertIs(result.client, client) + self.assertIsInstance(result, operation.Operation) + self.assertEqual(result.operation.name, OP_NAME) self.assertIsNone(result.metadata) - self.assertEqual(result.caller_metadata, - {'request_type': 'CreateCluster'}) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] @@ -278,7 +276,7 @@ def test_create(self): def test_update(self): import datetime from google.longrunning import operations_pb2 - from google.cloud.operation import Operation + from google.cloud.future import operation from google.protobuf.any_pb2 import Any from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.bigtable._generated import ( @@ -324,15 +322,11 @@ def test_update(self): result = cluster.update() - self.assertIsInstance(result, Operation) - self.assertEqual(result.name, OP_NAME) - self.assertIs(result.target, cluster) - self.assertIs(result.client, client) + self.assertIsInstance(result, operation.Operation) + self.assertEqual(result.operation.name, OP_NAME) self.assertIsInstance(result.metadata, messages_v2_pb2.UpdateClusterMetadata) self.assertEqual(result.metadata.request_time, NOW_PB) - self.assertEqual(result.caller_metadata, - {'request_type': 'UpdateCluster'}) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] @@ -448,6 +442,7 @@ class _Client(object): def __init__(self, project): self.project = project self.project_name = 'projects/' + self.project + self._operations_stub = mock.sentinel.operations_stub def __eq__(self, other): return (other.project == self.project and diff --git a/bigtable/tests/unit/test_instance.py b/bigtable/tests/unit/test_instance.py index cdad3c376d0a..03c0034fc49e 100644 --- a/bigtable/tests/unit/test_instance.py +++ b/bigtable/tests/unit/test_instance.py @@ -15,6 +15,8 @@ import unittest +import mock + class TestInstance(unittest.TestCase): @@ -236,7 +238,7 @@ def test_create(self): bigtable_instance_admin_pb2 as messages_v2_pb2) from google.cloud._helpers import _datetime_to_pb_timestamp from tests.unit._testing import _FakeStub - from google.cloud.operation import Operation + from google.cloud.future import operation from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES NOW = datetime.datetime.utcnow() @@ -263,15 +265,11 @@ def test_create(self): # Perform the method and check the result. result = instance.create() - self.assertIsInstance(result, Operation) - self.assertEqual(result.name, self.OP_NAME) - self.assertIs(result.target, instance) - self.assertIs(result.client, client) + self.assertIsInstance(result, operation.Operation) + self.assertEqual(result.operation.name, self.OP_NAME) self.assertIsInstance(result.metadata, messages_v2_pb2.CreateInstanceMetadata) self.assertEqual(result.metadata.request_time, NOW_PB) - self.assertEqual(result.caller_metadata, - {'request_type': 'CreateInstance'}) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] @@ -291,7 +289,7 @@ def test_create_w_explicit_serve_nodes(self): from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from tests.unit._testing import _FakeStub - from google.cloud.operation import Operation + from google.cloud.future import operation SERVE_NODES = 5 @@ -308,10 +306,8 @@ def test_create_w_explicit_serve_nodes(self): # Perform the method and check the result. result = instance.create() - self.assertIsInstance(result, Operation) - self.assertEqual(result.name, self.OP_NAME) - self.assertIs(result.target, instance) - self.assertIs(result.client, client) + self.assertIsInstance(result, operation.Operation) + self.assertEqual(result.operation.name, self.OP_NAME) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] @@ -582,6 +578,7 @@ class _Client(object): def __init__(self, project): self.project = project self.project_name = 'projects/' + self.project + self._operations_stub = mock.sentinel.operations_stub def copy(self): from copy import deepcopy