fix: usage of delete feature in tests

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,

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,
+                            dataset_name=getattr(args, "dataset_name", None),
-                            user_id=args.user_id,
+                            user_id=getattr(args, "user_id", None),
-                            all_data=args.all,
+                            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:
+                    elif hasattr(args, "dataset_name") or hasattr(args, "dataset_id"):
-                        dataset_id = 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

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_id)
-                user = await get_user(user_uuid)
             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

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

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
+        expected_user_id = uuid4()
-        mock_cognee = MagicMock()
+        expected_dataset_id = uuid4()
-        mock_cognee.delete = AsyncMock()
-        mock_get_deletion_counts = AsyncMock()
-        mock_get_deletion_counts.return_value = DeletionCountsPreview()
 
-        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+        get_user_mock.return_value = User(id=expected_user_id)
-            command = DeleteCommand()
-            args = argparse.Namespace(
-                dataset_name="test_dataset", user_id=None, all=False, force=False
-            )
 
-            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}'?")
+        command.execute(args)
-        assert mock_asyncio_run.call_count == 2
+
-        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        delete_dataset_mock.assert_awaited_once_with(
-        mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None)
+            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
+        command = DeleteCommand()
-        mock_cognee = MagicMock()
-        mock_cognee.delete = AsyncMock()
 
-        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+        expected_user_id = uuid4()
-            command = DeleteCommand()
-            args = argparse.Namespace(
-                dataset_name="test_dataset", user_id=None, all=False, force=True
-            )
 
-            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()
+        command.execute(args)
-        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+
-        mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None)
+        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)

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(
+    @patch("cognee.cli.commands.delete_command.fmt.confirm")
-        self, mock_asyncio_run, mock_confirm, mock_get_deletion_counts
+    @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
+        fmt_confirm_mock.return_value = True
-        mock_cognee = MagicMock()
+        delete_all_mock.return_value = None
-        mock_cognee.delete = AsyncMock()
-        mock_get_deletion_counts = AsyncMock()
-        mock_get_deletion_counts.return_value = DeletionCountsPreview()
 
-        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+        expected_user_id = uuid4()
-            command = DeleteCommand()
-            args = argparse.Namespace(dataset_name=None, user_id="test_user", all=True, force=False)
 
-            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?")
+        delete_all_mock.assert_called_once_with(user_id=expected_user_id)
-        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")
@@ -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()