Removing xrange and making test _Key return _Key instances.

googleapis · Dec 18, 2014 · 2ac3226 · 2ac3226
1 parent 483d05c
commit 2ac3226
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/gcloud/datastore/test_dataset.py b/gcloud/datastore/test_dataset.py
@@ -228,7 +228,7 @@ def test_allocate_ids(self):
         result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
 
         # Check the IDs returned match _PathElementProto.
-        self.assertEqual(result, range(NUM_IDS))
+        self.assertEqual([key._id for key in result], range(NUM_IDS))
 
         # Check connection is called correctly.
         self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
@@ -269,7 +269,7 @@ def allocate_ids(self, dataset_id, key_pbs):
         self._called_dataset_id = dataset_id
         self._called_key_pbs = key_pbs
         num_pbs = len(key_pbs)
-        return [_KeyProto(i) for i in xrange(num_pbs)]
+        return [_KeyProto(i) for i in range(num_pbs)]
 
 
 class _PathElementProto(object):

diff --git a/gcloud/datastore/test_entity.py b/gcloud/datastore/test_entity.py
@@ -239,10 +239,13 @@ class _Key(object):
     _key = 'KEY'
     _partial = False
     _path = None
+    _id = None
 
     def id(self, id_to_set):
         self._called_id = id_to_set
-        return id_to_set
+        clone = _Key()
+        clone._id = id_to_set
+        return clone
 
     def to_protobuf(self):
         return self._key