Merge pull request #483 from dhermes/fix-477-part6

Address sixth part of 477: Moves Dataset.get_entities to __init__ module.

Author: dhermes

gcloud/datastore/__init__.py

@@ -49,6 +49,8 @@
 from gcloud import credentials
 from gcloud.datastore import _implicit_environ
 from gcloud.datastore.connection import Connection
+from gcloud.datastore.dataset import Dataset
+from gcloud.datastore import helpers
 
 
 SCOPE = ('https://www.googleapis.com/auth/datastore ',
@@ -158,16 +160,57 @@ def _require_connection():
     return _implicit_environ.CONNECTION
 
 
-def get_entities(keys):
+def get_entities(keys, missing=None, deferred=None,
+                 connection=None, dataset_id=None):
     """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.
 
+    :type missing: an empty list or None.
+    :param missing: If a list is passed, the key-only entities returned
+                    by the backend as "missing" will be copied into it.
+                    Use only as a keyword param.
+
+    :type deferred: an empty list or None.
+    :param deferred: If a list is passed, the keys returned
+                     by the backend as "deferred" will be copied into it.
+                     Use only as a keyword param.
+
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: Optional. The connection used to connect to datastore.
+
+    :type dataset_id: :class:`str`.
+    :param dataset_id: Optional. The ID of the dataset.
+
     :rtype: list of :class:`gcloud.datastore.entity.Entity`
     :returns: The requested entities.
     """
-    return _require_dataset().get_entities(keys)
+    connection = connection or _require_connection()
+    dataset_id = dataset_id or _require_dataset().id()
+
+    entity_pbs = connection.lookup(
+        dataset_id=dataset_id,
+        key_pbs=[k.to_protobuf() for k in keys],
+        missing=missing, deferred=deferred,
+    )
+
+    new_dataset = Dataset(dataset_id, connection=connection)
+    if missing is not None:
+        missing[:] = [
+            helpers.entity_from_protobuf(missed_pb, dataset=new_dataset)
+            for missed_pb in missing]
+
+    if deferred is not None:
+        deferred[:] = [
+            helpers.key_from_protobuf(deferred_pb)
+            for deferred_pb in deferred]
+
+    entities = []
+    for entity_pb in entity_pbs:
+        entities.append(helpers.entity_from_protobuf(
+            entity_pb, dataset=new_dataset))
+    return entities
 
 
 def allocate_ids(incomplete_key, num_ids, connection=None, dataset_id=None):
@@ -180,10 +223,10 @@ def allocate_ids(incomplete_key, num_ids, connection=None, dataset_id=None):
     :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.
+    :param connection: Optional. The connection used to connect to datastore.
 
     :type dataset_id: :class:`str`.
-    :param dataset_id: Optional. The ID of the dataset used to allocate.
+    :param dataset_id: Optional. The ID of the dataset.
 
     :rtype: list of :class:`gcloud.datastore.key.Key`
     :returns: The (complete) keys allocated with `incomplete_key` as root.

gcloud/datastore/connection.py

@@ -14,13 +14,13 @@
 
 """Connections to gcloud datastore API servers."""
 
-from gcloud import connection as base_connection
+from gcloud import connection
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore import helpers
 from gcloud.datastore.dataset import Dataset
 
 
-class Connection(base_connection.Connection):
+class Connection(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=base_connection.Connection._EMPTY):
+    def transaction(self, transaction=connection.Connection._EMPTY):
         """Getter/setter for the connection's transaction object.
 
         :type transaction: :class:`gcloud.datastore.transaction.Transaction`,

gcloud/datastore/dataset.py

@@ -14,8 +14,6 @@
 
 """Create / interact with gcloud datastore datasets."""
 
-from gcloud.datastore import helpers
-
 
 class Dataset(object):
     """A dataset in the Cloud Datastore.
@@ -71,44 +69,3 @@ def id(self):
         """
 
         return self._id
-
-    def get_entities(self, keys, missing=None, deferred=None):
-        """Retrieves entities from the dataset, along with their attributes.
-
-        :type keys: list of :class:`gcloud.datastore.key.Key`
-        :param keys: List of keys to be retrieved.
-
-        :type missing: an empty list or None.
-        :param missing: If a list is passed, the key-only entities returned
-                        by the backend as "missing" will be copied into it.
-                        Use only as a keyword param.
-
-        :type deferred: an empty list or None.
-        :param deferred: If a list is passed, the keys returned
-                        by the backend as "deferred" will be copied into it.
-                        Use only as a keyword param.
-
-        :rtype: list of :class:`gcloud.datastore.entity.Entity`
-        :returns: The requested entities.
-        """
-        entity_pbs = self.connection().lookup(
-            dataset_id=self.id(),
-            key_pbs=[k.to_protobuf() for k in keys],
-            missing=missing, deferred=deferred,
-        )
-
-        if missing is not None:
-            missing[:] = [
-                helpers.entity_from_protobuf(missed_pb, dataset=self)
-                for missed_pb in missing]
-
-        if deferred is not None:
-            deferred[:] = [
-                helpers.key_from_protobuf(deferred_pb)
-                for deferred_pb in deferred]
-
-        entities = []
-        for entity_pb in entity_pbs:
-            entities.append(helpers.entity_from_protobuf(
-                entity_pb, dataset=self))
-        return entities

gcloud/datastore/key.py

@@ -238,12 +238,12 @@ def get(self, connection=None):
                   match found.
         """
         # Temporary import hack until Dataset is removed in #477.
-        from gcloud.datastore.dataset import Dataset
+        from gcloud import datastore
 
         # We allow partial keys to attempt a get, the backend will fail.
         connection = connection or _implicit_environ.CONNECTION
-        dataset = Dataset(self.dataset_id, connection=connection)
-        entities = dataset.get_entities([self])
+        entities = datastore.get_entities(
+            [self], connection=connection, dataset_id=self.dataset_id)
 
         if entities:
             result = entities[0]

gcloud/datastore/test___init__.py

@@ -189,21 +189,152 @@ 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
+
+class Test_get_entities_function(unittest2.TestCase):
+
+    def _callFUT(self, keys, missing=None, deferred=None,
+                 connection=None, dataset_id=None):
+        from gcloud.datastore import get_entities
+        return get_entities(keys, missing=missing, deferred=deferred,
+                            connection=connection, dataset_id=dataset_id)
+
+    def test_get_entities_miss(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        connection = _Connection()
+        key = Key('Kind', 1234, dataset_id=DATASET_ID)
+        results = self._callFUT([key], connection=connection,
+                                dataset_id=DATASET_ID)
+        self.assertEqual(results, [])
+
+    def test_get_entities_miss_w_missing(self):
+        from gcloud.datastore import datastore_v1_pb2 as datastore_pb
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        ID = 1234
+
+        # Make a missing entity pb to be returned from mock backend.
+        missed = datastore_pb.Entity()
+        missed.key.partition_id.dataset_id = DATASET_ID
+        path_element = missed.key.path_element.add()
+        path_element.kind = KIND
+        path_element.id = ID
+
+        # Set missing entity on mock connection.
+        connection = _Connection()
+        connection._missing = [missed]
+
+        key = Key(KIND, ID, dataset_id=DATASET_ID)
+        missing = []
+        entities = self._callFUT([key], connection=connection,
+                                 dataset_id=DATASET_ID, missing=missing)
+        self.assertEqual(entities, [])
+        self.assertEqual([missed.key().to_protobuf() for missed in missing],
+                         [key.to_protobuf()])
+
+    def test_get_entities_miss_w_deferred(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        key = Key('Kind', 1234, dataset_id=DATASET_ID)
+
+        # Set deferred entity on mock connection.
+        connection = _Connection()
+        connection._deferred = [key.to_protobuf()]
+
+        deferred = []
+        entities = self._callFUT([key], connection=connection,
+                                 dataset_id=DATASET_ID, deferred=deferred)
+        self.assertEqual(entities, [])
+        self.assertEqual([def_key.to_protobuf() for def_key in deferred],
+                         [key.to_protobuf()])
+
+    def _make_entity_pb(self, dataset_id, kind, integer_id, name, str_val):
+        from gcloud.datastore import datastore_v1_pb2 as datastore_pb
+
+        entity_pb = datastore_pb.Entity()
+        entity_pb.key.partition_id.dataset_id = dataset_id
+        path_element = entity_pb.key.path_element.add()
+        path_element.kind = kind
+        path_element.id = integer_id
+        prop = entity_pb.property.add()
+        prop.name = name
+        prop.value.string_value = str_val
+
+        return entity_pb
+
+    def test_get_entities_hit(self):
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
+
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+
+        # Make a found entity pb to be returned from mock backend.
+        entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID,
+                                         'foo', 'Foo')
+
+        # Make a connection to return the entity pb.
+        connection = _Connection(entity_pb)
+
+        key = Key(KIND, ID, dataset_id=DATASET_ID)
+        result, = self._callFUT([key], connection=connection,
+                                dataset_id=DATASET_ID)
+        new_key = result.key()
+
+        # Check the returned value is as expected.
+        self.assertFalse(new_key is key)
+        self.assertEqual(new_key.dataset_id, DATASET_ID)
+        self.assertEqual(new_key.path, PATH)
+        self.assertEqual(list(result), ['foo'])
+        self.assertEqual(result['foo'], 'Foo')
+
+    def test_get_entities_implicit(self):
         from gcloud.datastore import _implicit_environ
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_connection import _Connection
         from gcloud.datastore.test_entity import _Dataset
         from gcloud._testing import _Monkey
 
+        DATASET_ID = 'DATASET'
+        KIND = 'Kind'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+
+        # Make a found entity pb to be returned from mock backend.
+        entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID,
+                                         'foo', 'Foo')
+
+        # Make a connection to return the entity pb.
+        CUSTOM_CONNECTION = _Connection(entity_pb)
         CUSTOM_DATASET = _Dataset()
-        DUMMY_KEYS = [object(), object()]
-        DUMMY_VALS = [object(), object()]
-        for key, val in zip(DUMMY_KEYS, DUMMY_VALS):
-            CUSTOM_DATASET[key] = val
 
-        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
-            result = gcloud.datastore.get_entities(DUMMY_KEYS)
-        self.assertTrue(result == DUMMY_VALS)
+        key = Key(KIND, ID, dataset_id=DATASET_ID)
+        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET,
+                     CONNECTION=CUSTOM_CONNECTION):
+            result, = self._callFUT([key])
+
+        expected_called_with = {
+            'dataset_id': DATASET_ID,
+            'key_pbs': [key.to_protobuf()],
+        }
+        self.assertEqual(CUSTOM_CONNECTION._called_with, expected_called_with)
+
+        new_key = result.key()
+        # Check the returned value is as expected.
+        self.assertFalse(new_key is key)
+        self.assertEqual(new_key.dataset_id, DATASET_ID)
+        self.assertEqual(new_key.path, PATH)
+        self.assertEqual(list(result), ['foo'])
+        self.assertEqual(result['foo'], 'Foo')
 
 
 class Test_allocate_ids_function(unittest2.TestCase):
@@ -216,7 +347,7 @@ def _callFUT(self, incomplete_key, num_ids,
 
     def test_allocate_ids(self):
         from gcloud.datastore.key import Key
-        from gcloud.datastore.test_dataset import _Connection
+        from gcloud.datastore.test_connection import _Connection
 
         DATASET_ID = 'DATASET'
         INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
@@ -235,7 +366,7 @@ def test_allocate_ids(self):
     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_connection import _Connection
         from gcloud.datastore.test_entity import _Dataset
         from gcloud._testing import _Monkey
 
@@ -253,7 +384,7 @@ def test_allocate_ids_implicit(self):
     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_connection import _Connection
         from gcloud.datastore.test_entity import _Dataset
         from gcloud._testing import _Monkey