diff --git a/cognee/tests/unit/documents/pdf_document.py b/cognee/tests/unit/documents/pdf_document.py index ae4279729..7b51e2838 100644 --- a/cognee/tests/unit/documents/pdf_document.py +++ b/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"}, diff --git a/cognee/tests/unit/integration/test_model_to_graph_to_model.py b/cognee/tests/unit/integration/test_model_to_graph_to_model.py index d64f4e85a..e3ecaf0a0 100644 --- a/cognee/tests/unit/integration/test_model_to_graph_to_model.py +++ b/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", diff --git a/cognee/tests/unit/run_tasks/conftest.py b/cognee/tests/unit/run_tasks/conftest.py index 4d35a979b..6f7872657 100644 --- a/cognee/tests/unit/run_tasks/conftest.py +++ b/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 diff --git a/cognee/tests/unit/run_tasks/test_run_tasks.py b/cognee/tests/unit/run_tasks/test_run_tasks.py index 90e89ed5c..e76b8440b 100644 --- a/cognee/tests/unit/run_tasks/test_run_tasks.py +++ b/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(): 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 acd7b3c03..cf8047349 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 @@ -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 diff --git a/cognee/tests/unit/utils/get_mock_user.py b/cognee/tests/unit/utils/get_mock_user.py index ef8f000ac..76ce62cd0 100644 --- a/cognee/tests/unit/utils/get_mock_user.py +++ b/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()