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 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])