From ac5118ee34c4bd149ac26d042e2ffe5292ee3459 Mon Sep 17 00:00:00 2001
From: Andrej Milicevic <milicevi@Andrejs-MacBook-Pro.local>
Date: Wed, 15 Oct 2025 17:28:51 +0200
Subject: [PATCH] test:Add load test

---
 cognee/tests/load_test.py | 61 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)
 create mode 100644 cognee/tests/load_test.py

diff --git a/cognee/tests/load_test.py b/cognee/tests/load_test.py
new file mode 100644
index 000000000..da9b74ab9
--- /dev/null
+++ b/cognee/tests/load_test.py
@@ -0,0 +1,61 @@
+import os
+import pathlib
+import asyncio
+import time
+
+import cognee
+from cognee.modules.search.types import SearchType
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger()
+
+async def helper_func(num_of_searches):
+
+    start_time = time.time()
+
+    await cognee.cognify()
+
+    await asyncio.gather(
+        *[
+            cognee.search(query_text="Tell me about AI", query_type=SearchType.GRAPH_COMPLETION)
+            for _ in range(num_of_searches)
+        ]
+    )
+
+    end_time = time.time()
+
+    return end_time - start_time
+
+async def main():
+
+    file_path = os.path.join(
+        pathlib.Path(__file__).resolve().parent, "test_data/artificial-intelligence.pdf"
+    )
+
+    num_of_pdfs = 10
+    num_of_reps = 5
+    upper_boundary_minutes = 3
+    average_minutes = 1.5
+
+    await asyncio.gather(
+        *[
+            cognee.add(file_path, dataset_name=f"dataset_{i}")
+            for i in range(num_of_pdfs)
+        ]
+    )
+
+    recorded_times = await asyncio.gather(
+        *[helper_func(num_of_pdfs) for _ in range(num_of_reps)]
+    )
+
+    average_recorded_time = sum(recorded_times) / len(recorded_times)
+
+    assert average_recorded_time <= average_minutes * 60
+
+    assert all(rec_time <= upper_boundary_minutes * 60 for rec_time in recorded_times)
+
+    return
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
\ No newline at end of file