def check_DatastoreWriteFn(self, num_entities, use_fixed_batch_size=False):
"""A helper function to test _DatastoreWriteFn."""
with patch.object(helper, 'get_client',
return_value=self._mock_client):
entities = helper.create_entities(num_entities)
expected_entities = [
entity.to_client_entity() for entity in entities
]
# Infer project from write fn project arg.
if num_entities:
key = Key(['k1', 1234], project=self._PROJECT)
expected_key = key.to_client_key()
key.project = None
entities[0].key = key
expected_entities[0].key = expected_key
all_batch_entities = []
commit_count = [0]
self._mock_client.batch.side_effect = (lambda: FakeBatch(
all_batch_items=all_batch_entities, commit_count=commit_count))
datastore_write_fn = WriteToDatastore._DatastoreWriteFn(
self._PROJECT)
datastore_write_fn.start_bundle()
for entity in entities:
datastore_write_fn.process(entity)
datastore_write_fn.finish_bundle()
self.assertListEqual([e.key for e in all_batch_entities],
[e.key for e in expected_entities])
batch_count = math.ceil(num_entities / util.WRITE_BATCH_MAX_SIZE)
self.assertLessEqual(batch_count, commit_count[0])
def test_DatastoreDeleteFn(self):
with patch.object(helper, 'get_client',
return_value=self._mock_client):
keys = [entity.key for entity in helper.create_entities(10)]
expected_keys = [key.to_client_key() for key in keys]
# Infer project from delete fn project arg.
key = Key(['k1', 1234], project=self._PROJECT)
expected_key = key.to_client_key()
key.project = None
keys.append(key)
expected_keys.append(expected_key)
all_batch_keys = []
self._mock_client.batch.side_effect = (
lambda: FakeBatch(all_batch_items=all_batch_keys))
datastore_delete_fn = DeleteFromDatastore._DatastoreDeleteFn(
self._PROJECT)
datastore_delete_fn.start_bundle()
for key in keys:
datastore_delete_fn.process(key)
datastore_delete_fn.finish_bundle()
self.assertListEqual(all_batch_keys, expected_keys)
def test_DatastoreWriteLargeEntities(self):
"""100*100kB entities gets split over two Commit RPCs."""
with patch.object(helper, 'get_client', return_value=self._mock_client):
entities = helper.create_entities(100)
commit_count = [0]
self._mock_client.batch.side_effect = (
lambda: FakeBatch(commit_count=commit_count))
datastore_write_fn = WriteToDatastore._DatastoreWriteFn(self._PROJECT)
datastore_write_fn.start_bundle()
for entity in entities:
entity.set_properties({'large': u'A' * 100000})
datastore_write_fn.process(entity)
datastore_write_fn.finish_bundle()
self.assertEqual(2, commit_count[0])
