Moving connection.allocate_ids function to datastore.__init__.

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]

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]

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)

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