Implementing Bigtable Operation.finished().

dhermes · dhermes · commit 786720153618 · 2015-12-07T21:56:13.000-08:00
This method is intended to be used to check if a create,
update or undelete operation has completed.
diff --git a/gcloud/bigtable/cluster.py b/gcloud/bigtable/cluster.py
@@ -24,6 +24,7 @@
     bigtable_cluster_service_messages_pb2 as messages_pb2)
 from gcloud.bigtable._generated import (
     bigtable_table_service_messages_pb2 as table_messages_pb2)
+from gcloud.bigtable._generated import operations_pb2
 from gcloud.bigtable.table import Table
 
 
@@ -169,23 +170,54 @@ class Operation(object):
 
     :type begin: :class:`datetime.datetime`
     :param begin: The time when the operation was started.
+
+    :type cluster: :class:`Cluster`
+    :param cluster: The cluster that created the operation.
     """
 
-    def __init__(self, op_type, op_id, begin):
+    def __init__(self, op_type, op_id, begin, cluster=None):
         self.op_type = op_type
         self.op_id = op_id
         self.begin = begin
+        self._cluster = cluster
+        self._complete = False
 
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             return False
         return (other.op_type == self.op_type and
                 other.op_id == self.op_id and
-                other.begin == self.begin)
+                other.begin == self.begin and
+                other._cluster == self._cluster and
+                other._complete == self._complete)
 
     def __ne__(self, other):
         return not self.__eq__(other)
 
+    def finished(self):
+        """Check if the operation has finished.
+
+        :rtype: bool
+        :returns: A boolean indicating if the current operation has completed.
+        :raises: :class:`ValueError <exceptions.ValueError>` if the operation
+                 has already completed.
+        """
+        if self._complete:
+            raise ValueError('The operation has completed.')
+
+        operation_name = ('operations/' + self._cluster.name +
+                          '/operations/%d' % (self.op_id,))
+        request_pb = operations_pb2.GetOperationRequest(name=operation_name)
+        # We expact a `._generated.operations_pb2.Operation`.
+        operation_pb = self._cluster._client._operations_stub.GetOperation(
+            request_pb, self._cluster._client.timeout_seconds)
+
+        if operation_pb.done:
+            self._complete = True
+            return True
+        else:
+            return False
+
 
 class Cluster(object):
     """Representation of a Google Cloud Bigtable Cluster.
diff --git a/gcloud/bigtable/test_cluster.py b/gcloud/bigtable/test_cluster.py
@@ -25,16 +25,25 @@ def _getTargetClass(self):
     def _makeOne(self, *args, **kwargs):
         return self._getTargetClass()(*args, **kwargs)
 
-    def test_constructor(self):
+    def _constructor_test_helper(self, cluster=None):
         import datetime
         op_type = 'fake-op'
         op_id = 8915
         begin = datetime.datetime(2015, 10, 22, 1, 1)
-        operation = self._makeOne(op_type, op_id, begin)
+        operation = self._makeOne(op_type, op_id, begin, cluster=cluster)
 
         self.assertEqual(operation.op_type, op_type)
         self.assertEqual(operation.op_id, op_id)
         self.assertEqual(operation.begin, begin)
+        self.assertEqual(operation._cluster, cluster)
+        self.assertFalse(operation._complete)
+
+    def test_constructor_defaults(self):
+        self._constructor_test_helper()
+
+    def test_constructor_explicit_cluster(self):
+        cluster = object()
+        self._constructor_test_helper(cluster=cluster)
 
     def test___eq__(self):
         import datetime
@@ -65,6 +74,66 @@ def test___ne__(self):
         operation2 = self._makeOne('bar', 456, None)
         self.assertNotEqual(operation1, operation2)
 
+    def test_finished_without_operation(self):
+        operation = self._makeOne(None, None, None)
+        operation._complete = True
+        with self.assertRaises(ValueError):
+            operation.finished()
+
+    def _finished_helper(self, done):
+        import datetime
+        from gcloud.bigtable._generated import operations_pb2
+        from gcloud.bigtable._testing import _FakeStub
+        from gcloud.bigtable.cluster import Cluster
+
+        project = 'PROJECT'
+        zone = 'zone'
+        cluster_id = 'cluster-id'
+        op_type = 'fake-op'
+        op_id = 789
+        begin = datetime.datetime(2015, 10, 22, 1, 1)
+        timeout_seconds = 1
+
+        client = _Client(project, timeout_seconds=timeout_seconds)
+        cluster = Cluster(zone, cluster_id, client)
+        operation = self._makeOne(op_type, op_id, begin, cluster=cluster)
+
+        # Create request_pb
+        op_name = ('operations/projects/' + project + '/zones/' +
+                   zone + '/clusters/' + cluster_id +
+                   '/operations/%d' % (op_id,))
+        request_pb = operations_pb2.GetOperationRequest(name=op_name)
+
+        # Create response_pb
+        response_pb = operations_pb2.Operation(done=done)
+
+        # Patch the stub used by the API method.
+        client._operations_stub = stub = _FakeStub(response_pb)
+
+        # Create expected_result.
+        expected_result = done
+
+        # Perform the method and check the result.
+        result = operation.finished()
+
+        self.assertEqual(result, expected_result)
+        self.assertEqual(stub.method_calls, [(
+            'GetOperation',
+            (request_pb, timeout_seconds),
+            {},
+        )])
+
+        if done:
+            self.assertTrue(operation._complete)
+        else:
+            self.assertFalse(operation._complete)
+
+    def test_finished(self):
+        self._finished_helper(done=True)
+
+    def test_finished_not_done(self):
+        self._finished_helper(done=False)
+
 
 class TestCluster(unittest2.TestCase):
 
