From 77732d86981193c3f46d45ec76ecd7db573400be Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Tue, 14 Oct 2025 12:45:05 +0200 Subject: [PATCH] fix: usage of delete feature in tests --- cognee/api/v1/update/update.py | 2 +- cognee/cli/commands/delete_command.py | 47 +++++++++----- .../data/methods/get_deletion_counts.py | 18 +++--- cognee/modules/users/methods/get_user.py | 1 - .../cli_unit_tests/test_cli_commands.py | 62 +++++++++---------- .../cli_unit_tests/test_cli_edge_cases.py | 42 ++++++------- 6 files changed, 92 insertions(+), 80 deletions(-) diff --git a/cognee/api/v1/update/update.py b/cognee/api/v1/update/update.py index c0d1e4afb..f890b27ca 100644 --- a/cognee/api/v1/update/update.py +++ b/cognee/api/v1/update/update.py @@ -74,7 +74,7 @@ async def update( - Execution timestamps and metadata """ if not user: - user = get_default_user() + user = await get_default_user() await datasets.delete_data( dataset_id=dataset_id, diff --git a/cognee/cli/commands/delete_command.py b/cognee/cli/commands/delete_command.py index ea75d28a5..d23c362b4 100644 --- a/cognee/cli/commands/delete_command.py +++ b/cognee/cli/commands/delete_command.py @@ -38,22 +38,30 @@ Be careful with deletion operations as they are irreversible. import cognee # Validate arguments - if not any([args.dataset_name, args.dataset_id, args.data_id, args.user_id, args.all]): + if not any( + [ + hasattr(args, "dataset_name"), + hasattr(args, "dataset_id"), + hasattr(args, "data_id"), + hasattr(args, "user_id"), + hasattr(args, "all"), + ] + ): fmt.error( "Please specify what to delete: --dataset-name, --dataset-id, --data-id, --user-id, or --all" ) return # If --force is used, skip the preview and go straight to deletion - if not args.force: + if not getattr(args, "force", False): # --- START PREVIEW LOGIC --- fmt.echo("Gathering data for preview...") try: preview_data = asyncio.run( get_deletion_counts( - dataset_name=args.dataset_name, - user_id=args.user_id, - all_data=args.all, + dataset_name=getattr(args, "dataset_name", None), + user_id=getattr(args, "user_id", None), + all_data=getattr(args, "all", False), ) ) except CliCommandException as e: @@ -66,25 +74,25 @@ Be careful with deletion operations as they are irreversible. fmt.echo("You are about to delete:") fmt.echo( - f"Datasets: {preview_data.datasets}\nEntries: {preview_data.entries}\nUsers: {preview_data.users}" + f"Datasets: {preview_data.datasets}\nEntries: {preview_data.data_entries}\nUsers: {preview_data.users}" ) fmt.echo("-" * 20) # --- END PREVIEW LOGIC --- # Build operation message for success/failure logging - if args.all: + if getattr(args, "all", False): confirm_msg = "Delete ALL data from cognee?" operation = "all data" - elif args.dataset_name: + elif hasattr(args, "dataset_name"): confirm_msg = f"Delete dataset '{args.dataset_name}'?" operation = f"dataset '{args.dataset_name}'" - elif args.user_id: + elif hasattr(args, "user_id"): confirm_msg = f"Delete all data for user '{args.user_id}'?" operation = f"data for user '{args.user_id}'" else: operation = "data" - if not args.force: + if not getattr(args, "force", False): fmt.warning("This operation is irreversible!") if not fmt.confirm(confirm_msg): fmt.echo("Deletion cancelled.") @@ -95,12 +103,16 @@ Be careful with deletion operations as they are irreversible. # Run the async delete function async def run_delete(): try: - if args.all: + if getattr(args, "all", False): + if not hasattr(args, "user_id"): + raise CliCommandException( + "No user ID provided for '--all' deletion. Please specify using --user-id param." + ) await cognee.datasets.delete_all(user_id=args.user_id) - elif args.dataset_name or args.dataset_id: - dataset_id = args.dataset_id + elif hasattr(args, "dataset_name") or hasattr(args, "dataset_id"): + dataset_id = getattr(args, "dataset_id", None) - if args.dataset_name and not args.dataset_id: + if hasattr(args, "dataset_name") and not hasattr(args, "dataset_id"): datasets = await get_datasets_by_name( args.dataset_name, user_id=args.user_id ) @@ -113,10 +125,15 @@ Be careful with deletion operations as they are irreversible. dataset = datasets[0] dataset_id = dataset.id + if not hasattr(args, "user_id"): + raise CliCommandException( + "No user ID provided for deletion. Please specify using --user-id param." + ) + await cognee.datasets.delete_dataset( dataset_id=dataset_id, user_id=args.user_id ) - elif args.dataset_id and args.data_id: + elif hasattr(args, "dataset_id") and hasattr(args, "data_id"): await cognee.datasets.delete_data(args.dataset_id, args.data_id) except Exception as e: raise CliCommandInnerException(f"Failed to delete: {str(e)}") from e diff --git a/cognee/modules/data/methods/get_deletion_counts.py b/cognee/modules/data/methods/get_deletion_counts.py index 9237a8b85..1076f028b 100644 --- a/cognee/modules/data/methods/get_deletion_counts.py +++ b/cognee/modules/data/methods/get_deletion_counts.py @@ -1,3 +1,4 @@ +from typing import Optional from uuid import UUID from cognee.cli.exceptions import CliCommandException from cognee.infrastructure.databases.exceptions.exceptions import EntityNotFoundError @@ -18,7 +19,7 @@ class DeletionCountsPreview: async def get_deletion_counts( - dataset_name: str = None, user_id: str = None, all_data: bool = False + dataset_name: Optional[str] = None, user_id: Optional[UUID] = None, all_data: bool = False ) -> DeletionCountsPreview: """ Calculates the number of items that will be deleted based on the provided arguments. @@ -47,7 +48,7 @@ async def get_deletion_counts( data_entry_count = (await session.execute(count_query)).scalar_one() counts.users = 1 counts.datasets = 1 - counts.entries = data_entry_count + counts.data_entries = data_entry_count return counts elif all_data: @@ -55,7 +56,7 @@ async def get_deletion_counts( counts.datasets = ( await session.execute(select(func.count()).select_from(Dataset)) ).scalar_one() - counts.entries = ( + counts.data_entries = ( await session.execute(select(func.count()).select_from(Data)) ).scalar_one() counts.users = ( @@ -67,8 +68,7 @@ async def get_deletion_counts( elif user_id: user = None try: - user_uuid = UUID(user_id) - user = await get_user(user_uuid) + user = await get_user(user_id) except (ValueError, EntityNotFoundError): raise CliCommandException(f"No User exists with ID {user_id}", error_code=1) counts.users = 1 @@ -79,14 +79,16 @@ async def get_deletion_counts( counts.datasets = dataset_count if dataset_count > 0: dataset_ids = [d.id for d in user_datasets] - # Count all data entries across all of the user's datasets + # Count all data data_entries across all of the user's datasets data_count_query = ( select(func.count()) .select_from(DatasetData) .where(DatasetData.dataset_id.in_(dataset_ids)) ) data_entry_count = (await session.execute(data_count_query)).scalar_one() - counts.entries = data_entry_count + counts.data_entries = data_entry_count else: - counts.entries = 0 + counts.data_entries = 0 return counts + + return counts diff --git a/cognee/modules/users/methods/get_user.py b/cognee/modules/users/methods/get_user.py index 2678a5a01..6fe89c804 100644 --- a/cognee/modules/users/methods/get_user.py +++ b/cognee/modules/users/methods/get_user.py @@ -1,7 +1,6 @@ from uuid import UUID from sqlalchemy import select from sqlalchemy.orm import selectinload -import sqlalchemy.exc from cognee.infrastructure.databases.relational import get_relational_engine from cognee.infrastructure.databases.exceptions import EntityNotFoundError from ..models import User 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 7654a781a..3f825a64c 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 @@ -6,6 +6,7 @@ import pytest import sys import argparse import asyncio +from uuid import uuid4 from unittest.mock import patch, MagicMock, AsyncMock, ANY from cognee.cli.commands.add_command import AddCommand from cognee.cli.commands.search_command import SearchCommand @@ -14,6 +15,7 @@ from cognee.cli.commands.delete_command import DeleteCommand from cognee.cli.commands.config_command import ConfigCommand from cognee.cli.exceptions import CliCommandException from cognee.modules.data.methods.get_deletion_counts import DeletionCountsPreview +from cognee.modules.users.models import User # Mock asyncio.run to properly handle coroutines @@ -283,33 +285,29 @@ class TestDeleteCommand: assert "all" in actions assert "force" in actions - @patch("cognee.cli.commands.delete_command.get_deletion_counts") + @patch("cognee.api.v1.datasets.datasets.delete_dataset") + @patch("cognee.modules.data.methods.get_deletion_counts.get_user") @patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) def test_execute_delete_dataset_with_confirmation( - self, mock_asyncio_run, mock_confirm, mock_get_deletion_counts + self, mock_asyncio_run, mock_confirm, get_user_mock, delete_dataset_mock ): """Test execute delete dataset with user confirmation""" - # Mock the cognee module - mock_cognee = MagicMock() - mock_cognee.delete = AsyncMock() - mock_get_deletion_counts = AsyncMock() - mock_get_deletion_counts.return_value = DeletionCountsPreview() + expected_user_id = uuid4() + expected_dataset_id = uuid4() - with patch.dict(sys.modules, {"cognee": mock_cognee}): - command = DeleteCommand() - args = argparse.Namespace( - dataset_name="test_dataset", user_id=None, all=False, force=False - ) + get_user_mock.return_value = User(id=expected_user_id) - mock_confirm.return_value = True + command = DeleteCommand() + args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id) - command.execute(args) + mock_confirm.return_value = True - mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?") - assert mock_asyncio_run.call_count == 2 - assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0]) - mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None) + command.execute(args) + + delete_dataset_mock.assert_awaited_once_with( + dataset_id=expected_dataset_id, user_id=expected_user_id + ) @patch("cognee.cli.commands.delete_command.get_deletion_counts") @patch("cognee.cli.commands.delete_command.fmt.confirm") @@ -327,29 +325,29 @@ class TestDeleteCommand: mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?") + @patch("cognee.api.v1.datasets.datasets.delete_all") @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) - def test_execute_delete_forced(self, mock_asyncio_run): + def test_execute_delete_forced(self, mock_asyncio_run, delete_all_mock): """Test execute delete with force flag""" - # Mock the cognee module - mock_cognee = MagicMock() - mock_cognee.delete = AsyncMock() + command = DeleteCommand() - with patch.dict(sys.modules, {"cognee": mock_cognee}): - command = DeleteCommand() - args = argparse.Namespace( - dataset_name="test_dataset", user_id=None, all=False, force=True - ) + expected_user_id = uuid4() - command.execute(args) + args = argparse.Namespace( + dataset_name="test_dataset", + force=True, + all=True, + user_id=expected_user_id, + ) - mock_asyncio_run.assert_called_once() - assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0]) - mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None) + command.execute(args) + + delete_all_mock.assert_awaited_once_with(user_id=expected_user_id) def test_execute_no_delete_target(self): """Test execute when no delete target is specified""" command = DeleteCommand() - args = argparse.Namespace(dataset_name=None, user_id=None, all=False, force=False) + args = argparse.Namespace(force=True) # Should not raise exception, just return with error message command.execute(args) diff --git a/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py b/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py index ca27c0f67..344c61ed0 100644 --- a/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py +++ b/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py @@ -6,13 +6,14 @@ import pytest import sys import asyncio import argparse -from unittest.mock import patch, MagicMock, AsyncMock, ANY, call +from uuid import uuid4 +from unittest.mock import patch, MagicMock, AsyncMock, ANY from cognee.cli.commands.add_command import AddCommand from cognee.cli.commands.search_command import SearchCommand from cognee.cli.commands.cognify_command import CognifyCommand from cognee.cli.commands.delete_command import DeleteCommand from cognee.cli.commands.config_command import ConfigCommand -from cognee.cli.exceptions import CliCommandException, CliCommandInnerException +from cognee.cli.exceptions import CliCommandException from cognee.modules.data.methods.get_deletion_counts import DeletionCountsPreview @@ -379,32 +380,27 @@ class TestCognifyCommandEdgeCases: class TestDeleteCommandEdgeCases: """Test edge cases for DeleteCommand""" - @patch("cognee.cli.commands.delete_command.get_deletion_counts") - @patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) - def test_delete_all_with_user_id( - self, mock_asyncio_run, mock_confirm, mock_get_deletion_counts - ): + @patch("cognee.cli.commands.delete_command.fmt.confirm") + @patch("cognee.api.v1.datasets.datasets.delete_all") + def test_delete_all_with_user_id(self, delete_all_mock, fmt_confirm_mock, async_run_mock): """Test delete command with both --all and --user-id""" - # Mock the cognee module - mock_cognee = MagicMock() - mock_cognee.delete = AsyncMock() - mock_get_deletion_counts = AsyncMock() - mock_get_deletion_counts.return_value = DeletionCountsPreview() + fmt_confirm_mock.return_value = True + delete_all_mock.return_value = None - with patch.dict(sys.modules, {"cognee": mock_cognee}): - command = DeleteCommand() - args = argparse.Namespace(dataset_name=None, user_id="test_user", all=True, force=False) + expected_user_id = uuid4() - mock_confirm.return_value = True + command = DeleteCommand() + args = argparse.Namespace( + dataset_name=None, + user_id=expected_user_id, + all=True, + force=False, + ) - # Should handle both flags being set - command.execute(args) + command.execute(args) - mock_confirm.assert_called_once_with("Delete ALL data from cognee?") - assert mock_asyncio_run.call_count == 2 - assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0]) - mock_cognee.delete.assert_awaited_once_with(dataset_name=None, user_id="test_user") + delete_all_mock.assert_called_once_with(user_id=expected_user_id) @patch("cognee.cli.commands.delete_command.get_deletion_counts") @patch("cognee.cli.commands.delete_command.fmt.confirm") @@ -414,7 +410,7 @@ class TestDeleteCommandEdgeCases: mock_get_deletion_counts.return_value = DeletionCountsPreview() command = DeleteCommand() - args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=False) + args = argparse.Namespace(dataset_name="test_dataset") mock_confirm.side_effect = KeyboardInterrupt()