From 7d3657aef7419431192ec457550f5b0a2e6e72b6 Mon Sep 17 00:00:00 2001
From: Leon Luithlen <leon@topoteretes.com>
Date: Tue, 12 Nov 2024 11:32:14 +0100
Subject: [PATCH] Add fixture for mock user

---
 cognee/tests/unit/run_tasks/conftest.py        | 18 ++++++++++++++++++
 cognee/tests/unit/run_tasks/test_run_tasks.py  |  6 ++----
 .../run_tasks/test_run_tasks_from_queue.py     |  6 ------
 cognee/tests/unit/utils/get_mock_user.py       |  1 +
 4 files changed, 21 insertions(+), 10 deletions(-)
 create mode 100644 cognee/tests/unit/run_tasks/conftest.py

diff --git a/cognee/tests/unit/run_tasks/conftest.py b/cognee/tests/unit/run_tasks/conftest.py
new file mode 100644
index 000000000..4d35a979b
--- /dev/null
+++ b/cognee/tests/unit/run_tasks/conftest.py
@@ -0,0 +1,18 @@
+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
+
+
+@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
diff --git a/cognee/tests/unit/run_tasks/test_run_tasks.py b/cognee/tests/unit/run_tasks/test_run_tasks.py
index ff309a8a3..90e89ed5c 100644
--- a/cognee/tests/unit/run_tasks/test_run_tasks.py
+++ b/cognee/tests/unit/run_tasks/test_run_tasks.py
@@ -1,8 +1,10 @@
 import asyncio
+import warnings
 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
 
 
 async def run_and_check_tasks():
@@ -40,8 +42,4 @@ 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 5902090b3..acd7b3c03 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,8 +2,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
-from unittest.mock import patch
-from cognee.tests.unit.utils.get_mock_user import get_mock_user
 
 
 async def pipeline(data_queue):
@@ -51,8 +49,4 @@ 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
index e81f10d0a..ef8f000ac 100644
--- a/cognee/tests/unit/utils/get_mock_user.py
+++ b/cognee/tests/unit/utils/get_mock_user.py
@@ -1,6 +1,7 @@
 from unittest.mock import Mock, create_autospec
 from uuid import UUID, uuid4
 from typing import Optional, List
+from cognee.modules.users.models import User
 
 
 def get_mock_user(