diff --git a/cognee/tests/unit/run_tasks/test_run_tasks.py b/cognee/tests/unit/run_tasks/test_run_tasks.py index f3de0f068..ff309a8a3 100644 --- a/cognee/tests/unit/run_tasks/test_run_tasks.py +++ b/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()) diff --git a/cognee/tests/unit/run_tasks/test_run_tasks_from_queue.py b/cognee/tests/unit/run_tasks/test_run_tasks_from_queue.py index 7aaf20fcf..5902090b3 100644 --- a/cognee/tests/unit/run_tasks/test_run_tasks_from_queue.py +++ b/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()) diff --git a/cognee/tests/unit/utils/get_mock_user.py b/cognee/tests/unit/utils/get_mock_user.py new file mode 100644 index 000000000..e81f10d0a --- /dev/null +++ b/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