fix: change the way datasets import is mocked

This commit is contained in:
Boris Arzentar 2025-10-14 14:35:25 +02:00
parent f83fa12703
commit 344b057d44
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
2 changed files with 24 additions and 21 deletions

View file

@ -285,12 +285,11 @@ class TestDeleteCommand:
assert "all" in actions
assert "force" in actions
@patch("cognee.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, get_user_mock, delete_dataset_mock
self, mock_asyncio_run, mock_confirm, get_user_mock
):
"""Test execute delete dataset with user confirmation"""
expected_user_id = uuid4()
@ -298,16 +297,20 @@ class TestDeleteCommand:
get_user_mock.return_value = User(id=expected_user_id)
command = DeleteCommand()
args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
mock_cognee = MagicMock()
mock_cognee.datasets.delete_dataset = AsyncMock()
mock_confirm.return_value = True
with patch.dict(sys.modules, {"cognee": mock_cognee}):
command = DeleteCommand()
args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
command.execute(args)
mock_confirm.return_value = True
delete_dataset_mock.assert_awaited_once_with(
dataset_id=expected_dataset_id, user_id=expected_user_id
)
command.execute(args)
mock_cognee.datasets.delete_dataset.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")

View file

@ -382,25 +382,25 @@ class TestDeleteCommandEdgeCases:
@patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
@patch("cognee.cli.commands.delete_command.fmt.confirm")
@patch("cognee.datasets.delete_all")
def test_delete_all_with_user_id(self, delete_all_mock, fmt_confirm_mock, async_run_mock):
def test_delete_all_with_user_id(self, fmt_confirm_mock, async_run_mock):
"""Test delete command with both --all and --user-id"""
fmt_confirm_mock.return_value = True
delete_all_mock.return_value = None
expected_user_id = uuid4()
command = DeleteCommand()
args = argparse.Namespace(
dataset_name=None,
user_id=expected_user_id,
all=True,
force=False,
)
mock_cognee = MagicMock()
mock_cognee.datasets.delete_all = AsyncMock()
command.execute(args)
with patch.dict(sys.modules, {"cognee": mock_cognee}):
command = DeleteCommand()
args = argparse.Namespace(
user_id=expected_user_id,
all=True,
)
delete_all_mock.assert_called_once_with(user_id=expected_user_id)
command.execute(args)
mock_cognee.datasets.delete_all.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")