From 1864fc8f2b39dfba4abb644efa2a2794cf9e2ff9 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Thu, 30 Oct 2025 22:23:01 +0100 Subject: [PATCH] fix: unit tests --- .../modules/pipelines/operations/run_tasks.py | 4 ++-- .../operations/run_tasks_distributed.py | 5 ++--- .../cli_unit_tests/test_cli_commands.py | 19 ++++++++++++++----- cognee/tests/test_feedback_enrichment.py | 8 +++++++- 4 files changed, 25 insertions(+), 11 deletions(-) diff --git a/cognee/modules/pipelines/operations/run_tasks.py b/cognee/modules/pipelines/operations/run_tasks.py index 22e146b49..4414b73bb 100644 --- a/cognee/modules/pipelines/operations/run_tasks.py +++ b/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 diff --git a/cognee/modules/pipelines/operations/run_tasks_distributed.py b/cognee/modules/pipelines/operations/run_tasks_distributed.py index f659c14e7..997ffd571 100644 --- a/cognee/modules/pipelines/operations/run_tasks_distributed.py +++ b/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 diff --git a/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py b/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py index e496b025e..f7ac5423e 100644 --- a/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py +++ b/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()) diff --git a/cognee/tests/test_feedback_enrichment.py b/cognee/tests/test_feedback_enrichment.py index 02d90db32..501be7b22 100644 --- a/cognee/tests/test_feedback_enrichment.py +++ b/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", )