Adding back implicit allocate_ids() to datastore.__init__.

dhermes · dhermes · commit 26917a40b46b · 2015-01-06T09:59:14.000-08:00
diff --git 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
@@ -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
@@ -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()
