fix: unit tests

cognee/modules/pipelines/operations/run_tasks.py

@@ -7,7 +7,7 @@ from uuid import UUID
 
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.relational import get_relational_engine
-from cognee.modules.data.methods import get_dataset
+from cognee.modules.data.methods import get_authorized_dataset
 from cognee.modules.pipelines.operations.run_tasks_distributed import run_tasks_distributed
 from cognee.modules.users.models import User
 from cognee.shared.logging_utils import get_logger
@@ -65,7 +65,7 @@ async def run_tasks(
     if not user:
         user = await get_default_user()
 
-    dataset = await get_dataset(user.id, dataset_id)
+    dataset = await get_authorized_dataset(user, dataset_id, "write")
     pipeline_id = generate_pipeline_id(user.id, dataset.id, pipeline_name)
     pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data)
     pipeline_run_id = pipeline_run.pipeline_run_id

cognee/modules/pipelines/operations/run_tasks_distributed.py

@@ -6,10 +6,9 @@ except ModuleNotFoundError:
 from typing import Any, List, Optional
 from uuid import UUID
 
-from cognee.modules.data.methods import get_dataset
+from cognee.modules.data.methods import get_authorized_dataset
 from cognee.modules.data.models import Dataset
 from cognee.modules.pipelines.tasks.task import Task
-from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.pipelines.models import (
     PipelineRunStarted,
     PipelineRunCompleted,
@@ -94,7 +93,7 @@ async def run_tasks_distributed(
     if not user:
         user = await get_default_user()
 
-    dataset = await get_dataset(user.id, dataset_id)
+    dataset = await get_authorized_dataset(user, dataset_id, "write")
     pipeline_id: UUID = generate_pipeline_id(user.id, dataset.id, pipeline_name)
     pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data)
     pipeline_run_id: UUID = pipeline_run.pipeline_run_id

cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py

@@ -288,13 +288,15 @@ class TestDeleteCommand:
         assert "all" in actions
         assert "force" in actions
 
-    @patch("cognee.modules.data.methods.get_deletion_counts.get_user")
-    @patch("cognee.cli.commands.delete_command.cognee_datasets.delete_dataset")
+    # @patch("cognee.modules.users.methods.get_user")
+    @patch("cognee.cli.commands.delete_command.get_user")
+    @patch("cognee.cli.commands.delete_command.cognee_datasets")
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
+    @patch("cognee.cli.commands.delete_command.get_deletion_counts")
     @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
     # @pytest.mark.asyncio
     def test_execute_delete_dataset_with_confirmation(
-        self, mock_asyncio_run, mock_confirm, delete_dataset_mock, get_user_mock
+        self, mock_asyncio_run, get_deletion_counts_mock, mock_confirm, datasets_mock, get_user_mock
     ):
         """Test execute delete dataset with user confirmation"""
         data_directory_path = os.path.join(
@@ -314,8 +316,15 @@ class TestDeleteCommand:
 
         expected_user_id = uuid4()
         expected_dataset_id = uuid4()
+        expected_user = User(id=expected_user_id)
 
-        get_user_mock.return_value = User(id=expected_user_id)
+        get_user_mock.return_value = expected_user
+
+        get_deletion_counts_mock = AsyncMock()
+        get_deletion_counts_mock.return_value = DeletionCountsPreview()
+
+        delete_dataset_mock = AsyncMock()
+        datasets_mock.delete_dataset = delete_dataset_mock
 
         command = DeleteCommand()
         args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
@@ -325,7 +334,7 @@ class TestDeleteCommand:
         command.execute(args)
 
         delete_dataset_mock.assert_awaited_once_with(
-            dataset_id=expected_dataset_id, user_id=expected_user_id
+            dataset_id=expected_dataset_id, user=expected_user
         )
 
         asyncio.run(cognee.prune.prune_data())

cognee/tests/test_feedback_enrichment.py

@@ -11,8 +11,11 @@ Tests the complete feedback enrichment pipeline:
 
 import os
 import pathlib
+from uuid import UUID, uuid4
 from collections import Counter
 
+from pydantic import BaseModel
+
 import cognee
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.pipelines.tasks.task import Task
@@ -129,10 +132,13 @@ async def main():
         Task(link_enrichments_to_feedback),
     ]
 
+    class EnrichmentData(BaseModel):
+        id: UUID
+
     await cognee.memify(
         extraction_tasks=extraction_tasks,
         enrichment_tasks=enrichment_tasks,
-        data=[{}],
+        data=[EnrichmentData(id=uuid4())],
         dataset="feedback_enrichment_test_memify",
     )