fix: change the way datasets import is mocked
This commit is contained in:
parent
f83fa12703
commit
344b057d44
2 changed files with 24 additions and 21 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue