gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

test: Add prune step to the test

Browse source
This commit is contained in:
Andrej Milicevic 2025-10-15 17:58:05 +02:00
parent ac5118ee34
commit c16459d236
1 changed files with 5 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
cognee/tests/load_test.py
View file

@ -9,7 +9,7 @@ from cognee.shared.logging_utils import get_logger
logger = get_logger() logger = get_logger()
async def helper_func(num_of_searches): async def process_and_search(num_of_searches):
start_time = time.time() start_time = time.time()
@ -37,6 +37,9 @@ async def main():
upper_boundary_minutes = 3 upper_boundary_minutes = 3
average_minutes = 1.5 average_minutes = 1.5
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
await asyncio.gather( await asyncio.gather(
*[ *[
cognee.add(file_path, dataset_name=f"dataset_{i}") cognee.add(file_path, dataset_name=f"dataset_{i}")
@ -45,7 +48,7 @@ async def main():
) )
recorded_times = await asyncio.gather( recorded_times = await asyncio.gather(
*[helper_func(num_of_pdfs) for _ in range(num_of_reps)] *[process_and_search(num_of_pdfs) for _ in range(num_of_reps)]
) )
average_recorded_time = sum(recorded_times) / len(recorded_times) average_recorded_time = sum(recorded_times) / len(recorded_times)
@ -54,8 +57,6 @@ async def main():
assert all(rec_time <= upper_boundary_minutes * 60 for rec_time in recorded_times) assert all(rec_time <= upper_boundary_minutes * 60 for rec_time in recorded_times)
return
if __name__ == "__main__": if __name__ == "__main__":
asyncio.run(main()) asyncio.run(main())