From c401bbc43abdf82312ca2cba5711555aad45c866 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Fri, 2 Jan 2015 14:51:20 -0800
Subject: [PATCH 1/3] Moves Dataset.allocate_ids to connection module.

Addresses fifth part of #477.
---
 gcloud/datastore/__init__.py        | 19 ++-------------
 gcloud/datastore/connection.py      | 38 ++++++++++++++++++++++++++---
 gcloud/datastore/dataset.py         | 26 --------------------
 gcloud/datastore/test___init__.py   | 16 ------------
 gcloud/datastore/test_connection.py | 32 ++++++++++++++++++++++++
 gcloud/datastore/test_dataset.py    | 25 -------------------
 gcloud/datastore/test_entity.py     |  3 ---
 regression/datastore.py             |  5 +++-
 8 files changed, 73 insertions(+), 91 deletions(-)

diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py
index 7ab8342ae1fa..30bfbd23d2ad 100644
--- a/gcloud/datastore/__init__.py
+++ b/gcloud/datastore/__init__.py
@@ -48,7 +48,7 @@
 
 from gcloud import credentials
 from gcloud.datastore import _implicit_environ
-from gcloud.datastore.connection import Connection
+from gcloud.datastore import connection as connection_module
 
 
 SCOPE = ('https://www.googleapis.com/auth/datastore ',
@@ -104,7 +104,7 @@ def get_connection():
     """
     implicit_credentials = credentials.get_credentials()
     scoped_credentials = implicit_credentials.create_scoped(SCOPE)
-    return Connection(credentials=scoped_credentials)
+    return connection_module.Connection(credentials=scoped_credentials)
 
 
 def get_dataset(dataset_id):
@@ -156,18 +156,3 @@ def get_entities(keys):
     :returns: The requested entities.
     """
     return _require_dataset().get_entities(keys)
-
-
-def allocate_ids(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: The 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`
-    :returns: The (complete) keys allocated with `incomplete_key` as root.
-    """
-    return _require_dataset().allocate_ids(incomplete_key, num_ids)
diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py
index 13e676d7481f..250b2cb8388e 100644
--- a/gcloud/datastore/connection.py
+++ b/gcloud/datastore/connection.py
@@ -14,13 +14,13 @@
 
 """Connections to gcloud datastore API servers."""
 
-from gcloud import connection
+from gcloud import connection as base_connection
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore import helpers
 from gcloud.datastore.dataset import Dataset
 
 
-class Connection(connection.Connection):
+class Connection(base_connection.Connection):
     """A connection to the Google Cloud Datastore via the Protobuf API.
 
     This class should understand only the basic types (and protobufs)
@@ -125,7 +125,7 @@ def build_api_url(cls, dataset_id, method, base_url=None,
             api_version=(api_version or cls.API_VERSION),
             dataset_id=dataset_id, method=method)
 
-    def transaction(self, transaction=connection.Connection._EMPTY):
+    def transaction(self, transaction=base_connection.Connection._EMPTY):
         """Getter/setter for the connection's transaction object.
 
         :type transaction: :class:`gcloud.datastore.transaction.Transaction`,
@@ -575,3 +575,35 @@ def _copy_deferred_keys(lookup_request, lookup_response):
         lookup_request.key.remove(old_key)
     for def_key in lookup_response.deferred:
         lookup_request.key.add().CopyFrom(def_key)
+
+
+def allocate_ids(incomplete_key, num_ids, connection, dataset_id):
+    """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.
+
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: The connection used to allocate IDs.
+
+    :type dataset_id: :class:`str`.
+    :param dataset_id: The ID of the dataset used to allocate.
+
+    :rtype: list of :class:`gcloud.datastore.key.Key`
+    :returns: 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 = connection.allocate_ids(dataset_id, incomplete_key_pbs)
+    allocated_ids = [allocated_key_pb.path_element[-1].id
+                     for allocated_key_pb in allocated_key_pbs]
+    return [incomplete_key.completed_key(allocated_id)
+            for allocated_id in allocated_ids]
diff --git a/gcloud/datastore/dataset.py b/gcloud/datastore/dataset.py
index f0902c13c5e1..659c2b90d07e 100644
--- a/gcloud/datastore/dataset.py
+++ b/gcloud/datastore/dataset.py
@@ -112,29 +112,3 @@ 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`
-        :returns: 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.completed_key(allocated_id)
-                for allocated_id in allocated_ids]
diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py
index eab3686eca79..b9af30c85108 100644
--- a/gcloud/datastore/test___init__.py
+++ b/gcloud/datastore/test___init__.py
@@ -181,19 +181,3 @@ def test_get_entities(self):
         with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
             result = gcloud.datastore.get_entities(DUMMY_KEYS)
         self.assertTrue(result == DUMMY_VALS)
-
-    def test_allocate_ids(self):
-        import gcloud.datastore
-        from gcloud.datastore import _implicit_environ
-        from gcloud.datastore.key import Key
-        from gcloud.datastore.test_entity import _Dataset
-        from gcloud._testing import _Monkey
-
-        CUSTOM_DATASET = _Dataset()
-        NUM_IDS = 2
-        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
-            INCOMPLETE_KEY = Key('KIND')
-            result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
-
-        # Check the IDs returned.
-        self.assertEqual([key.id for key in result], range(1, NUM_IDS + 1))
diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py
index a61b29ddbb8f..3bc52f487556 100644
--- a/gcloud/datastore/test_connection.py
+++ b/gcloud/datastore/test_connection.py
@@ -1139,6 +1139,38 @@ def mutation(self):
         self.assertEqual(len(mutation.delete), 1)
 
 
+class Test_allocate_ids_function(unittest2.TestCase):
+
+    def _callFUT(self, incomplete_key, num_ids, connection, dataset_id):
+        from gcloud.datastore.connection import allocate_ids
+        return allocate_ids(incomplete_key, num_ids, connection, dataset_id)
+
+    def test_allocate_ids(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_dataset import _Connection
+
+        DATASET_ID = 'DATASET'
+        INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
+        CONNECTION = _Connection()
+        NUM_IDS = 2
+        result = self._callFUT(INCOMPLETE_KEY, NUM_IDS,
+                               CONNECTION, DATASET_ID)
+
+        # Check the IDs returned match.
+        self.assertEqual([key.id for key in 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)
+
+    def test_allocate_ids_with_complete(self):
+        from gcloud.datastore.test_entity import _Key
+
+        COMPLETE_KEY = _Key()
+        self.assertRaises(ValueError, self._callFUT,
+                          COMPLETE_KEY, 2, None, None)
+
+
 class Http(object):
 
     _called_with = None
diff --git a/gcloud/datastore/test_dataset.py b/gcloud/datastore/test_dataset.py
index 5924ef444515..d335905e27e3 100644
--- a/gcloud/datastore/test_dataset.py
+++ b/gcloud/datastore/test_dataset.py
@@ -108,31 +108,6 @@ def test_get_entities_hit(self):
         self.assertEqual(list(result), ['foo'])
         self.assertEqual(result['foo'], 'Foo')
 
-    def test_allocate_ids(self):
-        from gcloud.datastore.key import Key
-
-        DATASET_ID = 'DATASET'
-        INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
-        CONNECTION = _Connection()
-        NUM_IDS = 2
-        DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
-        result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
-
-        # Check the IDs returned match.
-        self.assertEqual([key.id for key in 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)
-
-    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
diff --git a/gcloud/datastore/test_entity.py b/gcloud/datastore/test_entity.py
index 98bcba287083..60e02e9efd29 100644
--- a/gcloud/datastore/test_entity.py
+++ b/gcloud/datastore/test_entity.py
@@ -274,9 +274,6 @@ def connection(self):
     def get_entities(self, keys):
         return [self.get(key) for key in keys]
 
-    def allocate_ids(self, incomplete_key, num_ids):
-        return [incomplete_key.completed_key(i + 1) for i in range(num_ids)]
-
 
 class _Connection(object):
     _transaction = _saved = _deleted = None
diff --git a/regression/datastore.py b/regression/datastore.py
index eadf4224643a..8c96db28229f 100644
--- a/regression/datastore.py
+++ b/regression/datastore.py
@@ -18,6 +18,7 @@
 import unittest2
 
 from gcloud import datastore
+from gcloud.datastore.connection import allocate_ids
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 from gcloud.datastore.query import Query
@@ -48,7 +49,9 @@ class TestDatastoreAllocateIDs(TestDatastore):
     def test_allocate_ids(self):
         incomplete_key = Key('Kind')
         num_ids = 10
-        allocated_keys = datastore.allocate_ids(incomplete_key, num_ids)
+        connection = datastore.get_connection()
+        allocated_keys = allocate_ids(incomplete_key, num_ids,
+                                      connection, DATASET_ID)
         self.assertEqual(len(allocated_keys), num_ids)
 
         unique_ids = set()

From 26917a40b46bbcecbb995c29a75f4ead0e838de5 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Fri, 2 Jan 2015 15:11:53 -0800
Subject: [PATCH 2/3] Adding back implicit allocate_ids() to
 datastore.__init__.

---
 gcloud/datastore/__init__.py      | 30 ++++++++++++++
 gcloud/datastore/test___init__.py | 66 +++++++++++++++++++++++++------
 regression/datastore.py           |  5 +--
 3 files changed, 85 insertions(+), 16 deletions(-)

diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py
index 30bfbd23d2ad..6e70a98fc617 100644
--- a/gcloud/datastore/__init__.py
+++ b/gcloud/datastore/__init__.py
@@ -146,6 +146,18 @@ def _require_dataset():
     return _implicit_environ.DATASET
 
 
+def _require_connection():
+    """Convenience method to ensure CONNECTION is set.
+
+    :rtype: :class:`gcloud.datastore.connection.Connection`
+    :returns: A connection based on the current environment.
+    :raises: :class:`EnvironmentError` if CONNECTION is not set.
+    """
+    if _implicit_environ.CONNECTION is None:
+        raise EnvironmentError('Connection could not be inferred.')
+    return _implicit_environ.CONNECTION
+
+
 def get_entities(keys):
     """Retrieves entities from implied dataset, along with their attributes.
 
@@ -156,3 +168,21 @@ def get_entities(keys):
     :returns: The requested entities.
     """
     return _require_dataset().get_entities(keys)
+
+
+def allocate_ids(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: The 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`
+    :returns: The (complete) keys allocated with `incomplete_key` as root.
+    """
+    dataset = _require_dataset()
+    connection = _require_connection()
+    return connection_module.allocate_ids(incomplete_key, num_ids,
+                                          connection, dataset.id())
diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py
index b9af30c85108..3660984a6cae 100644
--- a/gcloud/datastore/test___init__.py
+++ b/gcloud/datastore/test___init__.py
@@ -151,20 +151,43 @@ def test_it(self):
 
 class Test_implicit_behavior(unittest2.TestCase):
 
-    def test__require_dataset(self):
+    def test__require_dataset_value_unset(self):
         import gcloud.datastore
         from gcloud.datastore import _implicit_environ
-        original_dataset = _implicit_environ.DATASET
-
-        try:
-            _implicit_environ.DATASET = None
-            self.assertRaises(EnvironmentError,
-                              gcloud.datastore._require_dataset)
-            NEW_DATASET = object()
-            _implicit_environ.DATASET = NEW_DATASET
-            self.assertEqual(gcloud.datastore._require_dataset(), NEW_DATASET)
-        finally:
-            _implicit_environ.DATASET = original_dataset
+        from gcloud._testing import _Monkey
+
+        with _Monkey(_implicit_environ, DATASET=None):
+            with self.assertRaises(EnvironmentError):
+                gcloud.datastore._require_dataset()
+
+    def test__require_dataset_value_set(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud._testing import _Monkey
+
+        FAKE_DATASET = object()
+        with _Monkey(_implicit_environ, DATASET=FAKE_DATASET):
+            stored_dataset = gcloud.datastore._require_dataset()
+        self.assertTrue(stored_dataset is FAKE_DATASET)
+
+    def test__require_connection_value_unset(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud._testing import _Monkey
+
+        with _Monkey(_implicit_environ, CONNECTION=None):
+            with self.assertRaises(EnvironmentError):
+                gcloud.datastore._require_connection()
+
+    def test__require_connection_value_set(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud._testing import _Monkey
+
+        FAKE_CONNECTION = object()
+        with _Monkey(_implicit_environ, CONNECTION=FAKE_CONNECTION):
+            stored_connection = gcloud.datastore._require_connection()
+        self.assertTrue(stored_connection is FAKE_CONNECTION)
 
     def test_get_entities(self):
         import gcloud.datastore
@@ -181,3 +204,22 @@ def test_get_entities(self):
         with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
             result = gcloud.datastore.get_entities(DUMMY_KEYS)
         self.assertTrue(result == DUMMY_VALS)
+
+    def test_allocate_ids(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_dataset import _Connection
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud._testing import _Monkey
+
+        CUSTOM_DATASET = _Dataset()
+        CUSTOM_CONNECTION = _Connection()
+        NUM_IDS = 2
+        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET,
+                     CONNECTION=CUSTOM_CONNECTION):
+            INCOMPLETE_KEY = Key('KIND')
+            result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
+
+        # Check the IDs returned.
+        self.assertEqual([key.id for key in result], range(NUM_IDS))
diff --git a/regression/datastore.py b/regression/datastore.py
index 8c96db28229f..eadf4224643a 100644
--- a/regression/datastore.py
+++ b/regression/datastore.py
@@ -18,7 +18,6 @@
 import unittest2
 
 from gcloud import datastore
-from gcloud.datastore.connection import allocate_ids
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 from gcloud.datastore.query import Query
@@ -49,9 +48,7 @@ class TestDatastoreAllocateIDs(TestDatastore):
     def test_allocate_ids(self):
         incomplete_key = Key('Kind')
         num_ids = 10
-        connection = datastore.get_connection()
-        allocated_keys = allocate_ids(incomplete_key, num_ids,
-                                      connection, DATASET_ID)
+        allocated_keys = datastore.allocate_ids(incomplete_key, num_ids)
         self.assertEqual(len(allocated_keys), num_ids)
 
         unique_ids = set()

From 24098aa477006a56a70cae131fd5a7fe18310500 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Tue, 6 Jan 2015 10:32:10 -0800
Subject: [PATCH 3/3] Moving connection.allocate_ids function to
 datastore.__init__.

---
 gcloud/datastore/__init__.py        | 33 ++++++++++++++++-----
 gcloud/datastore/connection.py      | 32 --------------------
 gcloud/datastore/test___init__.py   | 45 +++++++++++++++++++++++++++--
 gcloud/datastore/test_connection.py | 32 --------------------
 4 files changed, 68 insertions(+), 74 deletions(-)

diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py
index 6e70a98fc617..dcc72276a1b0 100644
--- a/gcloud/datastore/__init__.py
+++ b/gcloud/datastore/__init__.py
@@ -48,7 +48,7 @@
 
 from gcloud import credentials
 from gcloud.datastore import _implicit_environ
-from gcloud.datastore import connection as connection_module
+from gcloud.datastore.connection import Connection
 
 
 SCOPE = ('https://www.googleapis.com/auth/datastore ',
@@ -104,7 +104,7 @@ def get_connection():
     """
     implicit_credentials = credentials.get_credentials()
     scoped_credentials = implicit_credentials.create_scoped(SCOPE)
-    return connection_module.Connection(credentials=scoped_credentials)
+    return Connection(credentials=scoped_credentials)
 
 
 def get_dataset(dataset_id):
@@ -170,19 +170,36 @@ def get_entities(keys):
     return _require_dataset().get_entities(keys)
 
 
-def allocate_ids(incomplete_key, num_ids):
+def allocate_ids(incomplete_key, num_ids, connection=None, dataset_id=None):
     """Allocates a list of IDs from a partial key.
 
     :type incomplete_key: A :class:`gcloud.datastore.key.Key`
-    :param incomplete_key: The partial key to use as base for allocated IDs.
+    :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.
 
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: Optional. The connection used to allocate IDs.
+
+    :type dataset_id: :class:`str`.
+    :param dataset_id: Optional. The ID of the dataset used to allocate.
+
     :rtype: list of :class:`gcloud.datastore.key.Key`
     :returns: The (complete) keys allocated with `incomplete_key` as root.
+    :raises: `ValueError` if `incomplete_key` is not a partial key.
     """
-    dataset = _require_dataset()
-    connection = _require_connection()
-    return connection_module.allocate_ids(incomplete_key, num_ids,
-                                          connection, dataset.id())
+    connection = connection or _require_connection()
+    dataset_id = dataset_id or _require_dataset().id()
+
+    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 = connection.allocate_ids(dataset_id, incomplete_key_pbs)
+    allocated_ids = [allocated_key_pb.path_element[-1].id
+                     for allocated_key_pb in allocated_key_pbs]
+    return [incomplete_key.completed_key(allocated_id)
+            for allocated_id in allocated_ids]
diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py
index 250b2cb8388e..e00170eac70f 100644
--- a/gcloud/datastore/connection.py
+++ b/gcloud/datastore/connection.py
@@ -575,35 +575,3 @@ def _copy_deferred_keys(lookup_request, lookup_response):
         lookup_request.key.remove(old_key)
     for def_key in lookup_response.deferred:
         lookup_request.key.add().CopyFrom(def_key)
-
-
-def allocate_ids(incomplete_key, num_ids, connection, dataset_id):
-    """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.
-
-    :type connection: :class:`gcloud.datastore.connection.Connection`
-    :param connection: The connection used to allocate IDs.
-
-    :type dataset_id: :class:`str`.
-    :param dataset_id: The ID of the dataset used to allocate.
-
-    :rtype: list of :class:`gcloud.datastore.key.Key`
-    :returns: 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 = connection.allocate_ids(dataset_id, incomplete_key_pbs)
-    allocated_ids = [allocated_key_pb.path_element[-1].id
-                     for allocated_key_pb in allocated_key_pbs]
-    return [incomplete_key.completed_key(allocated_id)
-            for allocated_id in allocated_ids]
diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py
index 3660984a6cae..7ed7717cf824 100644
--- a/gcloud/datastore/test___init__.py
+++ b/gcloud/datastore/test___init__.py
@@ -205,8 +205,34 @@ def test_get_entities(self):
             result = gcloud.datastore.get_entities(DUMMY_KEYS)
         self.assertTrue(result == DUMMY_VALS)
 
+
+class Test_allocate_ids_function(unittest2.TestCase):
+
+    def _callFUT(self, incomplete_key, num_ids,
+                 connection=None, dataset_id=None):
+        from gcloud.datastore import allocate_ids
+        return allocate_ids(incomplete_key, num_ids, connection=connection,
+                            dataset_id=dataset_id)
+
     def test_allocate_ids(self):
-        import gcloud.datastore
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_dataset import _Connection
+
+        DATASET_ID = 'DATASET'
+        INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
+        CONNECTION = _Connection()
+        NUM_IDS = 2
+        result = self._callFUT(INCOMPLETE_KEY, NUM_IDS,
+                               connection=CONNECTION, dataset_id=DATASET_ID)
+
+        # Check the IDs returned match.
+        self.assertEqual([key.id for key in 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)
+
+    def test_allocate_ids_implicit(self):
         from gcloud.datastore import _implicit_environ
         from gcloud.datastore.key import Key
         from gcloud.datastore.test_dataset import _Connection
@@ -219,7 +245,22 @@ def test_allocate_ids(self):
         with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET,
                      CONNECTION=CUSTOM_CONNECTION):
             INCOMPLETE_KEY = Key('KIND')
-            result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
+            result = self._callFUT(INCOMPLETE_KEY, NUM_IDS)
 
         # Check the IDs returned.
         self.assertEqual([key.id for key in result], range(NUM_IDS))
+
+    def test_allocate_ids_with_complete(self):
+        from gcloud.datastore import _implicit_environ
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_dataset import _Connection
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud._testing import _Monkey
+
+        CUSTOM_DATASET = _Dataset()
+        CUSTOM_CONNECTION = _Connection()
+        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET,
+                     CONNECTION=CUSTOM_CONNECTION):
+            COMPLETE_KEY = Key('KIND', 1234)
+            self.assertRaises(ValueError, self._callFUT,
+                              COMPLETE_KEY, 2)
diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py
index 3bc52f487556..a61b29ddbb8f 100644
--- a/gcloud/datastore/test_connection.py
+++ b/gcloud/datastore/test_connection.py
@@ -1139,38 +1139,6 @@ def mutation(self):
         self.assertEqual(len(mutation.delete), 1)
 
 
-class Test_allocate_ids_function(unittest2.TestCase):
-
-    def _callFUT(self, incomplete_key, num_ids, connection, dataset_id):
-        from gcloud.datastore.connection import allocate_ids
-        return allocate_ids(incomplete_key, num_ids, connection, dataset_id)
-
-    def test_allocate_ids(self):
-        from gcloud.datastore.key import Key
-        from gcloud.datastore.test_dataset import _Connection
-
-        DATASET_ID = 'DATASET'
-        INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
-        CONNECTION = _Connection()
-        NUM_IDS = 2
-        result = self._callFUT(INCOMPLETE_KEY, NUM_IDS,
-                               CONNECTION, DATASET_ID)
-
-        # Check the IDs returned match.
-        self.assertEqual([key.id for key in 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)
-
-    def test_allocate_ids_with_complete(self):
-        from gcloud.datastore.test_entity import _Key
-
-        COMPLETE_KEY = _Key()
-        self.assertRaises(ValueError, self._callFUT,
-                          COMPLETE_KEY, 2, None, None)
-
-
 class Http(object):
 
     _called_with = None