Adding back implicit get_entities() to datastore.__init__.

dhermes · dhermes · commit 6d2a5a79c2fd · 2015-01-02T16:41:42.000-08:00
diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py
@@ -158,6 +158,20 @@ def _require_connection():
     return _implicit_environ.CONNECTION
 
 
+def get_entities(keys):
+    """Retrieves entities from implied dataset, along with their attributes.
+
+    :type keys: list of :class:`gcloud.datastore.key.Key`
+    :param keys: The name of the item to retrieve.
+
+    :rtype: list of :class:`gcloud.datastore.entity.Entity`
+    :returns: The requested entities.
+    """
+    dataset = _require_dataset()
+    connection = _require_connection()
+    return connection_module.get_entities(keys, connection, dataset.id())
+
+
 def allocate_ids(incomplete_key, num_ids):
     """Allocates a list of IDs from a partial key.
 
diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py
@@ -189,6 +189,38 @@ def test__require_connection_value_set(self):
             stored_connection = gcloud.datastore._require_connection()
         self.assertTrue(stored_connection is FAKE_CONNECTION)
 
+    def test_get_entities(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud.datastore import connection as connection_module
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud._testing import _Monkey
+
+        _called_keys = []
+        _called_connection = []
+        _called_dataset_id = []
+        RETURN_VALUE = object()
+
+        def mock_get_entities(keys, connection, dataset_id):
+            _called_keys.append(keys)
+            _called_connection.append(connection)
+            _called_dataset_id.append(dataset_id)
+            return RETURN_VALUE
+
+        DUMMY_KEYS = object()
+        CUSTOM_CONNECTION = object()
+        CUSTOM_DATASET = _Dataset()
+        with _Monkey(connection_module, get_entities=mock_get_entities):
+            with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET,
+                         CONNECTION=CUSTOM_CONNECTION):
+                result = gcloud.datastore.get_entities(DUMMY_KEYS)
+
+        # Compare results.
+        self.assertEqual(result, RETURN_VALUE)
+        self.assertEqual(_called_keys, [DUMMY_KEYS])
+        self.assertEqual(_called_connection, [CUSTOM_CONNECTION])
+        self.assertEqual(_called_dataset_id, [CUSTOM_DATASET.id()])
+
     def test_allocate_ids(self):
         import gcloud.datastore
         from gcloud.datastore import _implicit_environ
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 get_entities
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 from gcloud.datastore.query import Query
@@ -139,8 +138,7 @@ def test_save_multiple(self):
             self.case_entities_to_delete.append(entity2)
 
         keys = [entity1.key(), entity2.key()]
-        connection = datastore.get_connection()
-        matches = get_entities(keys, connection, DATASET_ID)
+        matches = datastore.get_entities(keys)
         self.assertEqual(len(matches), 2)
 
     def test_empty_kind(self):
