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 "all" in actions
assert "force" in actions assert "force" in actions
@patch("cognee.datasets.delete_dataset")
@patch("cognee.modules.data.methods.get_deletion_counts.get_user") @patch("cognee.modules.data.methods.get_deletion_counts.get_user")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")
@patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
def test_execute_delete_dataset_with_confirmation( 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""" """Test execute delete dataset with user confirmation"""
expected_user_id = uuid4() expected_user_id = uuid4()
@ -298,16 +297,20 @@ class TestDeleteCommand:
get_user_mock.return_value = User(id=expected_user_id) get_user_mock.return_value = User(id=expected_user_id)
command = DeleteCommand() mock_cognee = MagicMock()
args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id) 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( command.execute(args)
dataset_id=expected_dataset_id, user_id=expected_user_id
) 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.get_deletion_counts")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @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.asyncio.run", side_effect=_mock_run)
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")
@patch("cognee.datasets.delete_all") def test_delete_all_with_user_id(self, fmt_confirm_mock, async_run_mock):
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""" """Test delete command with both --all and --user-id"""
fmt_confirm_mock.return_value = True fmt_confirm_mock.return_value = True
delete_all_mock.return_value = None
expected_user_id = uuid4() expected_user_id = uuid4()
command = DeleteCommand() mock_cognee = MagicMock()
args = argparse.Namespace( mock_cognee.datasets.delete_all = AsyncMock()
dataset_name=None,
user_id=expected_user_id,
all=True,
force=False,
)
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.get_deletion_counts")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")