From 0194f30ad9724d50ea9a8ac77a7a984b4d29823a Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Thu, 18 Dec 2014 00:36:54 -0800
Subject: [PATCH] Adding a non-protobuf allocate_ids method on Dataset.

Partially address #336.
---
 gcloud/datastore/dataset.py      | 26 ++++++++++++++++
 gcloud/datastore/test_dataset.py | 52 ++++++++++++++++++++++++++++++++
 gcloud/datastore/test_entity.py  |  4 +++
 regression/datastore.py          | 14 ++++-----
 4 files changed, 88 insertions(+), 8 deletions(-)

diff --git a/gcloud/datastore/dataset.py b/gcloud/datastore/dataset.py
index c527c34aab1fa..6760a742e5a7c 100644
--- a/gcloud/datastore/dataset.py
+++ b/gcloud/datastore/dataset.py
@@ -182,3 +182,29 @@ def get_entities(self, keys, missing=None, deferred=None):
             entities.append(helpers.entity_from_protobuf(
                 entity_pb, dataset=self))
         return entities
+
+    def allocate_ids(self, incomplete_key, num_ids):
+        """Allocates a list of IDs from a partial key.
+
+        :type incomplete_key: A :class:`gcloud.datastore.key.Key`
+        :param incomplete_key: Partial key to use as base for allocated IDs.
+
+        :type num_ids: A :class:`int`.
+        :param num_ids: The number of IDs to allocate.
+
+        :rtype: list of :class:`gcloud.datastore.key.Key`
+        :return: The (complete) keys allocated with `incomplete_key` as root.
+        :raises: `ValueError` if `incomplete_key` is not a partial key.
+        """
+        if not incomplete_key.is_partial():
+            raise ValueError(('Key is not partial.', incomplete_key))
+
+        incomplete_key_pb = incomplete_key.to_protobuf()
+        incomplete_key_pbs = [incomplete_key_pb] * num_ids
+
+        allocated_key_pbs = self.connection().allocate_ids(
+            self.id(), incomplete_key_pbs)
+        allocated_ids = [allocated_key_pb.path_element[-1].id
+                         for allocated_key_pb in allocated_key_pbs]
+        return [incomplete_key.id(allocated_id)
+                for allocated_id in allocated_ids]
diff --git a/gcloud/datastore/test_dataset.py b/gcloud/datastore/test_dataset.py
index a59e4b307b178..f155198f963ee 100644
--- a/gcloud/datastore/test_dataset.py
+++ b/gcloud/datastore/test_dataset.py
@@ -213,6 +213,40 @@ def test_get_entities_hit(self):
         self.assertEqual(list(result), ['foo'])
         self.assertEqual(result['foo'], 'Foo')
 
+    def test_allocate_ids(self):
+        from gcloud.datastore.test_entity import _Key
+
+        INCOMPLETE_KEY = _Key()
+        PROTO_ID = object()
+        INCOMPLETE_KEY._key = _KeyProto(PROTO_ID)
+        INCOMPLETE_KEY._partial = True
+
+        CONNECTION = _Connection()
+        NUM_IDS = 2
+        DATASET_ID = 'foo'
+        DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
+        result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
+
+        # Check the IDs returned match _PathElementProto.
+        self.assertEqual(result, range(NUM_IDS))
+
+        # Check connection is called correctly.
+        self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
+        self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS)
+
+        # Check the IDs passed to Connection.allocate_ids.
+        key_paths = [key_pb.path_element[-1].id
+                     for key_pb in CONNECTION._called_key_pbs]
+        self.assertEqual(key_paths, [PROTO_ID] * NUM_IDS)
+
+    def test_allocate_ids_with_complete(self):
+        from gcloud.datastore.test_entity import _Key
+
+        COMPLETE_KEY = _Key()
+        DATASET = self._makeOne(None)
+        self.assertRaises(ValueError, DATASET.allocate_ids,
+                          COMPLETE_KEY, 2)
+
 
 class _Connection(object):
     _called_with = None
@@ -230,3 +264,21 @@ def lookup(self, **kw):
         if deferred is not None:
             deferred.extend(self._deferred)
         return self._result
+
+    def allocate_ids(self, dataset_id, key_pbs):
+        self._called_dataset_id = dataset_id
+        self._called_key_pbs = key_pbs
+        num_pbs = len(key_pbs)
+        return [_KeyProto(i) for i in xrange(num_pbs)]
+
+
+class _PathElementProto(object):
+
+    def __init__(self, _id):
+        self.id = _id
+
+
+class _KeyProto(object):
+
+    def __init__(self, id_):
+        self.path_element = [_PathElementProto(id_)]
diff --git a/gcloud/datastore/test_entity.py b/gcloud/datastore/test_entity.py
index 164a69aeb003d..b2cafe26485d3 100644
--- a/gcloud/datastore/test_entity.py
+++ b/gcloud/datastore/test_entity.py
@@ -240,6 +240,10 @@ class _Key(object):
     _partial = False
     _path = None
 
+    def id(self, id_to_set):
+        self._called_id = id_to_set
+        return id_to_set
+
     def to_protobuf(self):
         return self._key
 
diff --git a/regression/datastore.py b/regression/datastore.py
index 811e5ac6b24ee..c8670ae32892e 100644
--- a/regression/datastore.py
+++ b/regression/datastore.py
@@ -42,18 +42,16 @@ class TestDatastoreAllocateIDs(TestDatastore):
 
     def test_allocate_ids(self):
         incomplete_key = datastore.key.Key(path=[{'kind': 'Kind'}])
-        incomplete_key_pb = incomplete_key.to_protobuf()
-        incomplete_key_pbs = [incomplete_key_pb] * 10
-
-        connection = self.dataset.connection()
-        allocated_key_pbs = connection.allocate_ids(self.dataset.id(),
-                                                    incomplete_key_pbs)
-        allocated_keys = [datastore.helpers.key_from_protobuf(key_pb)
-                          for key_pb in allocated_key_pbs]
+        allocated_keys = self.dataset.allocate_ids(incomplete_key, 10)
         self.assertEqual(len(allocated_keys), 10)
+
+        unique_ids = set()
         for key in allocated_keys:
+            unique_ids.add(key.id())
             self.assertFalse(key.is_partial())
 
+        self.assertEqual(len(unique_ids), 10)
+
 
 class TestDatastoreSave(TestDatastore):