Implementing Cluster.create().

Author: dhermes

gcloud/bigtable/cluster.py

@@ -17,6 +17,7 @@
 
 import re
 
+from gcloud.bigtable._generated import bigtable_cluster_data_pb2 as data_pb2
 from gcloud.bigtable._generated import (
     bigtable_cluster_service_messages_pb2 as messages_pb2)
 from gcloud.bigtable.table import Table
@@ -52,6 +53,27 @@ def _get_pb_property_value(message_pb, property_name):
     return getattr(message_pb, property_name)
 
 
+def _prepare_create_request(cluster):
+    """Creates a protobuf request for a CreateCluster request.
+
+    :type cluster: :class:`Cluster`
+    :param cluster: The cluster to be created.
+
+    :rtype: :class:`.messages_pb2.CreateClusterRequest`
+    :returns: The CreateCluster request object containing the cluster info.
+    """
+    zone_full_name = ('projects/' + cluster._client.project +
+                      '/zones/' + cluster.zone)
+    return messages_pb2.CreateClusterRequest(
+        name=zone_full_name,
+        cluster_id=cluster.cluster_id,
+        cluster=data_pb2.Cluster(
+            display_name=cluster.display_name,
+            serve_nodes=cluster.serve_nodes,
+        ),
+    )
+
+
 class Cluster(object):
     """Representation of a Google Cloud Bigtable Cluster.
 
@@ -83,6 +105,7 @@ def __init__(self, zone, cluster_id, client,
         self.display_name = display_name or cluster_id
         self.serve_nodes = serve_nodes
         self._client = client
+        self._operation = None
 
     def table(self, table_id):
         """Factory to create a table associated with this cluster.
@@ -171,3 +194,27 @@ def reload(self):
         # NOTE: _update_from_pb does not check that the project, zone and
         #       cluster ID on the response match the request.
         self._update_from_pb(cluster_pb)
+
+    def create(self):
+        """Create this cluster.
+
+        .. note::
+
+            Uses the ``project``, ``zone`` and ``cluster_id`` on the current
+            :class:`Cluster` in addition to the ``display_name`` and
+            ``serve_nodes``. If you'd like to change them before creating,
+            reset the values via
+
+            .. code:: python
+
+                cluster.display_name = 'New display name'
+                cluster.cluster_id = 'i-changed-my-mind'
+
+            before calling :meth:`create`.
+        """
+        request_pb = _prepare_create_request(self)
+        # We expect an `operations_pb2.Operation`.
+        cluster_pb = self._client._cluster_stub.CreateCluster(
+            request_pb, self._client.timeout_seconds)
+
+        self._operation = cluster_pb.current_operation

gcloud/bigtable/test_cluster.py

@@ -171,7 +171,8 @@ def test_reload(self):
         zone = 'zone'
         cluster_id = 'cluster-id'
         timeout_seconds = 123
-        client = _Client(project=project, timeout_seconds=timeout_seconds)
+
+        client = _Client(project, timeout_seconds=timeout_seconds)
         cluster = self._makeOne(zone, cluster_id, client)
 
         # Create request_pb
@@ -210,6 +211,58 @@ def test_reload(self):
         self.assertEqual(cluster.serve_nodes, serve_nodes)
         self.assertEqual(cluster.display_name, display_name)
 
+    def test_create(self):
+        from gcloud._testing import _Monkey
+        from gcloud.bigtable._generated import (
+            bigtable_cluster_data_pb2 as data_pb2)
+        from gcloud.bigtable._generated import operations_pb2
+        from gcloud.bigtable._testing import _FakeStub
+        from gcloud.bigtable import cluster as MUT
+
+        project = 'PROJECT'
+        zone = 'zone'
+        cluster_id = 'cluster-id'
+        timeout_seconds = 578
+
+        client = _Client(project, timeout_seconds=timeout_seconds)
+        cluster = self._makeOne(zone, cluster_id, client)
+
+        # Create request_pb. Just a mock since we monkey patch
+        # _prepare_create_request
+        request_pb = object()
+
+        # Create response_pb
+        op_id = 5678
+        op_name = ('operations/projects/%s/zones/%s/clusters/%s/'
+                   'operations/%d' % (project, zone, cluster_id, op_id))
+        current_op = operations_pb2.Operation(name=op_name)
+        response_pb = data_pb2.Cluster(current_operation=current_op)
+
+        # Patch the stub used by the API method.
+        client._cluster_stub = stub = _FakeStub(response_pb)
+
+        # Create expected_result.
+        expected_result = None  # create() has no return value.
+
+        # Perform the method and check the result.
+        prep_create_called = []
+
+        def mock_prep_create_req(cluster):
+            prep_create_called.append(cluster)
+            return request_pb
+
+        with _Monkey(MUT, _prepare_create_request=mock_prep_create_req):
+            result = cluster.create()
+
+        self.assertEqual(result, expected_result)
+        self.assertEqual(stub.method_calls, [(
+            'CreateCluster',
+            (request_pb, timeout_seconds),
+            {},
+        )])
+        self.assertEqual(cluster._operation, current_op)
+        self.assertEqual(prep_create_called, [cluster])
+
 
 class Test__get_pb_property_value(unittest2.TestCase):
 
@@ -233,6 +286,39 @@ def test_with_value_unset_on_pb(self):
             self._callFUT(cluster_pb, 'serve_nodes')
 
 
+class Test__prepare_create_request(unittest2.TestCase):
+
+    def _callFUT(self, cluster):
+        from gcloud.bigtable.cluster import _prepare_create_request
+        return _prepare_create_request(cluster)
+
+    def test_it(self):
+        from gcloud.bigtable._generated import (
+            bigtable_cluster_data_pb2 as data_pb2)
+        from gcloud.bigtable._generated import (
+            bigtable_cluster_service_messages_pb2 as messages_pb2)
+        from gcloud.bigtable.cluster import Cluster
+
+        project = 'PROJECT'
+        zone = 'zone'
+        cluster_id = 'cluster-id'
+        display_name = u'DISPLAY_NAME'
+        serve_nodes = 8
+        client = _Client(project)
+
+        cluster = Cluster(zone, cluster_id, client,
+                          display_name=display_name, serve_nodes=serve_nodes)
+        request_pb = self._callFUT(cluster)
+        self.assertTrue(isinstance(request_pb,
+                                   messages_pb2.CreateClusterRequest))
+        self.assertEqual(request_pb.cluster_id, cluster_id)
+        self.assertEqual(request_pb.name,
+                         'projects/' + project + '/zones/' + zone)
+        self.assertTrue(isinstance(request_pb.cluster, data_pb2.Cluster))
+        self.assertEqual(request_pb.cluster.display_name, display_name)
+        self.assertEqual(request_pb.cluster.serve_nodes, serve_nodes)
+
+
 class _Client(object):
 
     def __init__(self, project, timeout_seconds=None):