gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: usage of delete feature in tests

Browse source
This commit is contained in:
Boris Arzentar 2025-10-14 12:45:05 +02:00
parent 4d1cf5393c
commit 77732d8698
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
6 changed files with 92 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/api/v1/update/update.py
View file

@ -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,

47
cognee/cli/commands/delete_command.py
View file

@ -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

18
cognee/modules/data/methods/get_deletion_counts.py
View file

@ -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

1
cognee/modules/users/methods/get_user.py
View file

@ -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

62
cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py
View file

@ -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)

42
cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py
View file

@ -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()