Add get_mock_user

cognee/tests/unit/run_tasks/test_run_tasks.py

@@ -1,6 +1,8 @@
 import asyncio
 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
 
 
 async def run_and_check_tasks():
@@ -26,6 +28,7 @@ async def run_and_check_tasks():
             Task(add_one_single),
         ],
         10,
+        pipeline_name="test_run_tasks",
     )
 
     results = [5, 7, 9, 11, 13, 15, 17, 19, 21, 23]
@@ -37,4 +40,8 @@ async def run_and_check_tasks():
 
 
 def test_run_tasks():
+    @patch("from cognee.modules.users.methods import get_default_user")
+    def get_mock_user_wrapper():
+        return get_mock_user(None, None)
+
     asyncio.run(run_and_check_tasks())

cognee/tests/unit/run_tasks/test_run_tasks_from_queue.py

@@ -2,6 +2,8 @@ import asyncio
 from queue import Queue
 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
 
 
 async def pipeline(data_queue):
@@ -23,7 +25,8 @@ async def pipeline(data_queue):
             Task(queue_consumer),
             Task(add_one),
             Task(multiply_by_two),
-        ]
+        ],
+        pipeline_name="test_run_tasks_from_queue",
     )
 
     results = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
@@ -48,4 +51,8 @@ async def run_queue():
 
 
 def test_run_tasks_from_queue():
+    @patch("from cognee.modules.users.methods import get_default_user")
+    def get_mock_user_wrapper():
+        return get_mock_user(None, None)
+
     asyncio.run(run_queue())

cognee/tests/unit/utils/get_mock_user.py

@@ -0,0 +1,40 @@
+from unittest.mock import Mock, create_autospec
+from uuid import UUID, uuid4
+from typing import Optional, List
+
+
+def get_mock_user(
+    user_id: Optional[UUID] = None,
+    groups: Optional[List["Group"]] = None,
+    **additional_attributes
+) -> Mock:
+    """
+    Creates a mock User instance with configurable attributes.
+
+    Args:
+        user_id: Optional UUID for the user. Generates random UUID if not provided.
+        groups: Optional list of group mocks to associate with the user.
+        **additional_attributes: Any additional attributes to set on the mock user.
+
+    Returns:
+        Mock: A configured mock User instance.
+    """
+    # Generate a random UUID if none provided
+    user_id = user_id or uuid4()
+
+    # Create base mock
+    mock_user = create_autospec(User, instance=True)
+
+    # Configure basic attributes
+    mock_user.id = user_id
+    mock_user.__tablename__ = "users"
+    mock_user.groups = groups or []
+
+    # Set polymorphic identity
+    mock_user.__mapper_args__ = {"polymorphic_identity": "user"}
+
+    # Add any additional attributes
+    for attr_name, attr_value in additional_attributes.items():
+        setattr(mock_user, attr_name, attr_value)
+
+    return mock_user