diff --git a/gcloud/datastore/test_dataset.py b/gcloud/datastore/test_dataset.py index f155198f963ee..955eebc8eeca4 100644 --- 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 index b2cafe26485d3..c036720d42417 100644 --- 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