From a92f4bdf3fd13ba55a1fd811fb2fbe4da40917cf Mon Sep 17 00:00:00 2001 From: Daulet Amirkhanov Date: Thu, 2 Oct 2025 15:05:39 +0100 Subject: [PATCH] fix: update failing tests and refactor delete_preview implementation --- cognee/cli/commands/delete_command.py | 29 +++++++---- .../data/methods/get_deletion_counts.py | 48 ++++++++----------- .../cli_unit_tests/test_cli_commands.py | 17 +++++-- .../cli_unit_tests/test_cli_edge_cases.py | 9 +++- 4 files changed, 60 insertions(+), 43 deletions(-) diff --git a/cognee/cli/commands/delete_command.py b/cognee/cli/commands/delete_command.py index 0139dbe1a..46cf41367 100644 --- a/cognee/cli/commands/delete_command.py +++ b/cognee/cli/commands/delete_command.py @@ -46,13 +46,17 @@ Be careful with deletion operations as they are irreversible. if not args.force: # --- START PREVIEW LOGIC --- fmt.echo("Gathering data for preview...") - preview_data = asyncio.run( - get_deletion_counts( - dataset_name=args.dataset_name, - user_id=args.user_id, - all_data=args.all, + try: + preview_data = asyncio.run( + get_deletion_counts( + dataset_name=args.dataset_name, + user_id=args.user_id, + all_data=args.all, + ) ) - ) + except CliCommandException as e: + fmt.error(f"Error occured when fetching preview data: {str(e)}") + return if not preview_data: fmt.success("No data found to delete.") @@ -63,22 +67,27 @@ Be careful with deletion operations as they are irreversible. f"Datasets: {preview_data.datasets}\nEntries: {preview_data.entries}\nUsers: {preview_data.users}" ) fmt.echo("-" * 20) - fmt.warning("This operation is irreversible!") - if not fmt.confirm("Proceed?"): - fmt.echo("Deletion cancelled.") - return # --- END PREVIEW LOGIC --- # Build operation message for success/failure logging if args.all: + confirm_msg = "Delete ALL data from cognee?" operation = "all data" elif args.dataset_name: + confirm_msg = f"Delete dataset '{args.dataset_name}'?" operation = f"dataset '{args.dataset_name}'" elif 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: + fmt.warning("This operation is irreversible!") + if not fmt.confirm(confirm_msg): + fmt.echo("Deletion cancelled.") + return + fmt.echo(f"Deleting {operation}...") # Run the async delete function diff --git a/cognee/modules/data/methods/get_deletion_counts.py b/cognee/modules/data/methods/get_deletion_counts.py index f985cf71d..9237a8b85 100644 --- a/cognee/modules/data/methods/get_deletion_counts.py +++ b/cognee/modules/data/methods/get_deletion_counts.py @@ -5,10 +5,9 @@ from sqlalchemy import select from sqlalchemy.sql import func from cognee.infrastructure.databases.relational import get_relational_engine from cognee.modules.data.models import Dataset, Data, DatasetData -from cognee.modules.users.models import User, DatasetDatabase -from cognee.modules.users.methods import get_user, get_default_user +from cognee.modules.users.models import User +from cognee.modules.users.methods import get_user from dataclasses import dataclass -import cognee.cli.echo as fmt @dataclass @@ -35,7 +34,6 @@ async def get_deletion_counts( dataset = dataset_result.scalar_one_or_none() if dataset is None: - fmt.error(f"No dataset with this name: {dataset_name}") raise CliCommandException( f"No Dataset exists with the name {dataset_name}", error_code=1 ) @@ -72,27 +70,23 @@ async def get_deletion_counts( user_uuid = UUID(user_id) user = await get_user(user_uuid) except (ValueError, EntityNotFoundError): - # Handles cases where user_id is not a valid UUID or user is not found - fmt.error(f"No user exists with ID {user_id}") raise CliCommandException(f"No User exists with ID {user_id}", error_code=1) - user = await get_user(user_uuid) - if user: - counts.users = 1 - # Find all datasets owned by this user - datasets_query = select(Dataset).where(Dataset.owner_id == user.id) - user_datasets = (await session.execute(datasets_query)).scalars().all() - dataset_count = len(user_datasets) - 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 - 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 - else: - counts.entries = 0 - return counts + counts.users = 1 + # Find all datasets owned by this user + datasets_query = select(Dataset).where(Dataset.owner_id == user.id) + user_datasets = (await session.execute(datasets_query)).scalars().all() + dataset_count = len(user_datasets) + 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 + 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 + else: + counts.entries = 0 + return counts 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 e5d6a0c64..7654a781a 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 @@ -12,7 +12,8 @@ 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 # Mock asyncio.run to properly handle coroutines @@ -282,13 +283,18 @@ class TestDeleteCommand: assert "all" in actions assert "force" in actions + @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_execute_delete_dataset_with_confirmation(self, mock_asyncio_run, mock_confirm): + def test_execute_delete_dataset_with_confirmation( + self, mock_asyncio_run, mock_confirm, mock_get_deletion_counts + ): """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() with patch.dict(sys.modules, {"cognee": mock_cognee}): command = DeleteCommand() @@ -301,13 +307,16 @@ class TestDeleteCommand: command.execute(args) mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?") - mock_asyncio_run.assert_called_once() + 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) + @patch("cognee.cli.commands.delete_command.get_deletion_counts") @patch("cognee.cli.commands.delete_command.fmt.confirm") - def test_execute_delete_cancelled(self, mock_confirm): + def test_execute_delete_cancelled(self, mock_confirm, mock_get_deletion_counts): """Test execute when user cancels deletion""" + mock_get_deletion_counts = AsyncMock() + mock_get_deletion_counts.return_value = DeletionCountsPreview() command = DeleteCommand() args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=False) 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 d4e4b0ee5..de8681be9 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 @@ -13,6 +13,7 @@ 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.modules.data.methods.get_deletion_counts import DeletionCountsPreview # Mock asyncio.run to properly handle coroutines @@ -396,13 +397,17 @@ class TestDeleteCommandEdgeCases: command.execute(args) mock_confirm.assert_called_once_with("Delete ALL data from cognee?") - mock_asyncio_run.assert_called_once() + 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") + @patch("cognee.cli.commands.delete_command.get_deletion_counts") @patch("cognee.cli.commands.delete_command.fmt.confirm") - def test_delete_confirmation_keyboard_interrupt(self, mock_confirm): + def test_delete_confirmation_keyboard_interrupt(self, mock_confirm, mock_get_deletion_counts): """Test delete command when user interrupts confirmation""" + mock_get_deletion_counts = AsyncMock() + mock_get_deletion_counts.return_value = DeletionCountsPreview() + command = DeleteCommand() args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=False)