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 "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")
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue