diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py index b8ef47907055..58eef968cbd3 100644 --- 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 index 6116dc0bf1b7..1a674f036474 100644 --- 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 index 1c2ff5063fe4..9ff2b2630a2f 100644 --- 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):