test: Add load test.

.github/workflows/e2e_tests.yml

@@ -330,4 +330,33 @@ jobs:
           DB_PORT: 5432
           DB_USERNAME: cognee
           DB_PASSWORD: cognee
-        run: uv run python ./cognee/tests/test_concurrent_subprocess_access.py
+        run: uv run python ./cognee/tests/test_concurrent_subprocess_access.py
+
+  test-load:
+    name: Test Load
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Cognee Setup
+        uses: ./.github/actions/cognee_setup
+        with:
+          python-version: '3.11.x'
+
+      - name: Dependencies already installed
+        run: echo "Dependencies already installed in setup"
+
+      - name: Run Load Test
+        env:
+          ENV: 'dev'
+          ENABLE_BACKEND_ACCESS_CONTROL: True
+          LLM_MODEL: ${{ secrets.LLM_MODEL }}
+          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
+          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
+          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
+          EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
+          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
+          EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
+        run: uv run python ./cognee/tests/test_load.py

cognee/tests/test_load.py

@@ -9,8 +9,8 @@ from cognee.shared.logging_utils import get_logger
 
 logger = get_logger()
 
-async def process_and_search(num_of_searches):
 
+async def process_and_search(num_of_searches):
     start_time = time.time()
 
     await cognee.cognify()
@@ -26,26 +26,34 @@ async def process_and_search(num_of_searches):
 
     return end_time - start_time
 
-async def main():
 
+async def main():
     file_path = os.path.join(
         pathlib.Path(__file__).resolve().parent, "test_data/artificial-intelligence.pdf"
     )
+    data_directory_path = str(
+        pathlib.Path(
+            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_load")
+        ).resolve()
+    )
+    cognee.config.data_root_directory(data_directory_path)
+    cognee_directory_path = str(
+        pathlib.Path(
+            os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_load")
+        ).resolve()
+    )
+    cognee.config.system_root_directory(cognee_directory_path)
 
     num_of_pdfs = 10
     num_of_reps = 5
-    upper_boundary_minutes = 3
+    upper_boundary_minutes = 10
-    average_minutes = 1.5
+    average_minutes = 8
 
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
 
-    await asyncio.gather(
+    for i in range(num_of_pdfs):
-        *[
+        await cognee.add(file_path, dataset_name=f"dataset_{i}")
-            cognee.add(file_path, dataset_name=f"dataset_{i}")
-            for i in range(num_of_pdfs)
-        ]
-    )
 
     recorded_times = await asyncio.gather(
         *[process_and_search(num_of_pdfs) for _ in range(num_of_reps)]
@@ -59,4 +67,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main())