Address fifth part of 477: Moves Dataset.allocate_ids to connection module.

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.
 
@@ -158,16 +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.
     """
-    return _require_dataset().allocate_ids(incomplete_key, num_ids)
+    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

@@ -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`,

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]

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
@@ -182,18 +205,62 @@ 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
         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):
+        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(1, NUM_IDS + 1))
+        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_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

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