Use patch

cognee/tests/unit/documents/pdf_document.py

@@ -1,6 +1,8 @@
 import os
 import uuid
-from cognee.modules.data.processing.document_types.PdfDocument import PdfDocument
+
+from cognee.modules.data.processing.document_types.PdfDocument import \
+    PdfDocument
 
 GROUND_TRUTH = [
     {"word_count": 879, "len_text": 5622, "cut_type": "sentence_end"},

cognee/tests/unit/integration/test_model_to_graph_to_model.py

@@ -1,13 +1,12 @@
-import pytest
-from enum import Enum
 from datetime import datetime, timezone
+from enum import Enum
 from typing import Optional
-from cognee.infrastructure.engine import DataPoint
-from cognee.modules.graph.utils import (
-    get_graph_from_model,
-    get_model_instance_from_graph,
-)
 
+import pytest
+
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.graph.utils import (get_graph_from_model,
+                                        get_model_instance_from_graph)
 
 EDGE_GROUND_TRUTH = (
     "boris",

cognee/tests/unit/run_tasks/conftest.py

@@ -1,18 +1,18 @@
+from unittest.mock import patch
+
 import pytest
-import warnings
-from cognee.tests.unit.utils.get_mock_user import get_mock_user
+
 from cognee.modules.users.methods.get_default_user import get_default_user
-import cognee
+from cognee.tests.unit.utils.get_mock_user import get_mock_user
 
 
 @pytest.fixture(autouse=True, scope="session")
 def set_get_mock_user_wrapper():
 
     def get_mock_user_wrapper():
-        warnings.warn("\n\n\n---------------USING MOCK USER--------------------\n\n\n")
         return get_mock_user(None, None)
 
-    get_default_user = cognee.modules.users.methods.get_default_user
-    cognee.modules.users.methods.get_default_user = get_mock_user_wrapper()
-    yield
-    cognee.modules.users.methods.get_default_user = get_default_user
+    with patch(
+        "cognee.modules.users.methods.get_default_user", get_mock_user_wrapper()
+    ):
+        yield

cognee/tests/unit/run_tasks/test_run_tasks.py

@@ -1,10 +1,11 @@
 import asyncio
 import warnings
+from unittest.mock import patch
+
 from cognee.modules.pipelines.operations.run_tasks import run_tasks
 from cognee.modules.pipelines.tasks.Task import Task
-from unittest.mock import patch
-from cognee.tests.unit.utils.get_mock_user import get_mock_user
 from cognee.modules.users.methods.get_default_user import get_default_user
+from cognee.tests.unit.utils.get_mock_user import get_mock_user
 
 
 async def run_and_check_tasks():

cognee/tests/unit/run_tasks/test_run_tasks_from_queue.py

@@ -1,5 +1,6 @@
 import asyncio
 from queue import Queue
+
 from cognee.modules.pipelines.operations.run_tasks import run_tasks
 from cognee.modules.pipelines.tasks.Task import Task
 

cognee/tests/unit/utils/get_mock_user.py

@@ -1,6 +1,8 @@
+import warnings
+from typing import List, Optional
 from unittest.mock import Mock, create_autospec
 from uuid import UUID, uuid4
-from typing import Optional, List
+
 from cognee.modules.users.models import User
 
 
@@ -20,6 +22,8 @@ def get_mock_user(
     Returns:
         Mock: A configured mock User instance.
     """
+    warnings.warn("\n\n\n---------------USING MOCK USER--------------------\n\n\n")
+
     # Generate a random UUID if none provided
     user_id = user_id or uuid4()