483 lines
17 KiB
Python
483 lines
17 KiB
Python
"""
|
||
Tests for individual CLI commands with proper mocking and coroutine handling.
|
||
"""
|
||
|
||
import pytest
|
||
import sys
|
||
import argparse
|
||
import asyncio
|
||
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
|
||
|
||
|
||
# Mock asyncio.run to properly handle coroutines
|
||
def _mock_run(coro):
|
||
# Create an event loop and run the coroutine
|
||
loop = asyncio.new_event_loop()
|
||
try:
|
||
return loop.run_until_complete(coro)
|
||
finally:
|
||
loop.close()
|
||
|
||
|
||
class TestAddCommand:
|
||
"""Test the AddCommand class"""
|
||
|
||
def test_command_properties(self):
|
||
"""Test basic command properties"""
|
||
command = AddCommand()
|
||
assert command.command_string == "add"
|
||
assert "Add data" in command.help_string
|
||
assert command.docs_url is not None
|
||
|
||
def test_configure_parser(self):
|
||
"""Test parser configuration"""
|
||
command = AddCommand()
|
||
parser = argparse.ArgumentParser()
|
||
|
||
command.configure_parser(parser)
|
||
|
||
# Check that required arguments are added
|
||
actions = {action.dest: action for action in parser._actions}
|
||
assert "data" in actions
|
||
assert "dataset_name" in actions
|
||
|
||
# Check data argument accepts multiple values
|
||
assert actions["data"].nargs == "+"
|
||
|
||
@patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run)
|
||
def test_execute_single_item(self, mock_asyncio_run):
|
||
"""Test execute with single data item"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.add = AsyncMock()
|
||
|
||
with patch.dict(sys.modules, {"cognee": mock_cognee}):
|
||
command = AddCommand()
|
||
args = argparse.Namespace(data=["test.txt"], dataset_name="test_dataset")
|
||
command.execute(args)
|
||
|
||
mock_asyncio_run.assert_called_once()
|
||
assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
|
||
mock_cognee.add.assert_awaited_once_with(data="test.txt", dataset_name="test_dataset")
|
||
|
||
@patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run)
|
||
def test_execute_multiple_items(self, mock_asyncio_run):
|
||
"""Test execute with multiple data items"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.add = AsyncMock()
|
||
|
||
with patch.dict(sys.modules, {"cognee": mock_cognee}):
|
||
command = AddCommand()
|
||
args = argparse.Namespace(data=["test1.txt", "test2.txt"], dataset_name="test_dataset")
|
||
command.execute(args)
|
||
|
||
mock_asyncio_run.assert_called_once()
|
||
assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
|
||
mock_cognee.add.assert_awaited_once_with(
|
||
data=["test1.txt", "test2.txt"], dataset_name="test_dataset"
|
||
)
|
||
|
||
@patch("cognee.cli.commands.add_command.asyncio.run")
|
||
def test_execute_with_exception(self, mock_asyncio_run):
|
||
"""Test execute handles exceptions properly"""
|
||
command = AddCommand()
|
||
args = argparse.Namespace(data=["test.txt"], dataset_name="test_dataset")
|
||
|
||
mock_asyncio_run.side_effect = Exception("Test error")
|
||
|
||
with pytest.raises(CliCommandException):
|
||
command.execute(args)
|
||
|
||
|
||
class TestSearchCommand:
|
||
"""Test the SearchCommand class"""
|
||
|
||
def test_command_properties(self):
|
||
"""Test basic command properties"""
|
||
command = SearchCommand()
|
||
assert command.command_string == "search"
|
||
assert "Search and query" in command.help_string
|
||
assert command.docs_url is not None
|
||
|
||
def test_configure_parser(self):
|
||
"""Test parser configuration"""
|
||
command = SearchCommand()
|
||
parser = argparse.ArgumentParser()
|
||
|
||
command.configure_parser(parser)
|
||
|
||
# Check that required arguments are added
|
||
actions = {action.dest: action for action in parser._actions}
|
||
assert "query_text" in actions
|
||
assert "query_type" in actions
|
||
assert "datasets" in actions
|
||
assert "top_k" in actions
|
||
assert "output_format" in actions
|
||
|
||
# Check default values
|
||
assert actions["query_type"].default == "GRAPH_COMPLETION"
|
||
assert actions["top_k"].default == 10
|
||
assert actions["output_format"].default == "pretty"
|
||
|
||
@patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run)
|
||
def test_execute_basic_search(self, mock_asyncio_run):
|
||
"""Test execute with basic search"""
|
||
# Mock the cognee module and SearchType
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.search = AsyncMock(return_value=["result1", "result2"])
|
||
mock_search_type = MagicMock()
|
||
mock_search_type.__getitem__.return_value = "GRAPH_COMPLETION"
|
||
|
||
with patch.dict(sys.modules, {"cognee": mock_cognee}):
|
||
command = SearchCommand()
|
||
args = argparse.Namespace(
|
||
query_text="test query",
|
||
query_type="GRAPH_COMPLETION",
|
||
datasets=None,
|
||
top_k=10,
|
||
system_prompt=None,
|
||
output_format="pretty",
|
||
)
|
||
command.execute(args)
|
||
|
||
mock_asyncio_run.assert_called_once()
|
||
assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
|
||
mock_cognee.search.assert_awaited_once_with(
|
||
query_text="test query",
|
||
query_type=ANY,
|
||
datasets=None,
|
||
top_k=10,
|
||
system_prompt_path="answer_simple_question.txt",
|
||
)
|
||
# verify the enum’s name separately
|
||
called_enum = mock_cognee.search.await_args.kwargs["query_type"]
|
||
assert called_enum.name == "GRAPH_COMPLETION"
|
||
|
||
@patch("cognee.cli.commands.search_command.asyncio.run")
|
||
def test_execute_with_exception(self, mock_asyncio_run):
|
||
"""Test execute handles exceptions properly"""
|
||
command = SearchCommand()
|
||
args = argparse.Namespace(
|
||
query_text="test query",
|
||
query_type="GRAPH_COMPLETION",
|
||
datasets=None,
|
||
top_k=10,
|
||
system_prompt=None,
|
||
output_format="pretty",
|
||
)
|
||
|
||
mock_asyncio_run.side_effect = Exception("Search error")
|
||
|
||
with pytest.raises(CliCommandException):
|
||
command.execute(args)
|
||
|
||
|
||
class TestCognifyCommand:
|
||
"""Test the CognifyCommand class"""
|
||
|
||
def test_command_properties(self):
|
||
"""Test basic command properties"""
|
||
command = CognifyCommand()
|
||
assert command.command_string == "cognify"
|
||
assert "Transform ingested data" in command.help_string
|
||
assert command.docs_url is not None
|
||
|
||
def test_configure_parser(self):
|
||
"""Test parser configuration"""
|
||
command = CognifyCommand()
|
||
parser = argparse.ArgumentParser()
|
||
|
||
command.configure_parser(parser)
|
||
|
||
# Check that arguments are added
|
||
actions = {action.dest: action for action in parser._actions}
|
||
assert "datasets" in actions
|
||
assert "chunk_size" in actions
|
||
assert "ontology_file" in actions
|
||
assert "chunker" in actions
|
||
assert "background" in actions
|
||
assert "verbose" in actions
|
||
|
||
# Check default values
|
||
assert actions["chunker"].default == "TextChunker"
|
||
|
||
@patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run)
|
||
def test_execute_basic_cognify(self, mock_asyncio_run):
|
||
"""Test execute with basic cognify"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.cognify = AsyncMock(return_value="success")
|
||
|
||
with patch.dict(sys.modules, {"cognee": mock_cognee}):
|
||
command = CognifyCommand()
|
||
args = argparse.Namespace(
|
||
datasets=None,
|
||
chunk_size=None,
|
||
ontology_file=None,
|
||
chunker="TextChunker",
|
||
background=False,
|
||
verbose=False,
|
||
)
|
||
command.execute(args)
|
||
|
||
mock_asyncio_run.assert_called_once()
|
||
assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
|
||
from cognee.modules.chunking import TextChunker
|
||
|
||
mock_cognee.cognify.assert_awaited_once_with(
|
||
datasets=None,
|
||
chunk_size=None,
|
||
ontology_file_path=None,
|
||
chunker=TextChunker,
|
||
run_in_background=False,
|
||
)
|
||
|
||
@patch("cognee.cli.commands.cognify_command.asyncio.run")
|
||
def test_execute_with_exception(self, mock_asyncio_run):
|
||
"""Test execute handles exceptions properly"""
|
||
command = CognifyCommand()
|
||
args = argparse.Namespace(
|
||
datasets=None,
|
||
chunk_size=None,
|
||
ontology_file=None,
|
||
chunker="TextChunker",
|
||
background=False,
|
||
verbose=False,
|
||
)
|
||
|
||
mock_asyncio_run.side_effect = Exception("Cognify error")
|
||
|
||
with pytest.raises(CliCommandException):
|
||
command.execute(args)
|
||
|
||
|
||
class TestDeleteCommand:
|
||
"""Test the DeleteCommand class"""
|
||
|
||
def test_command_properties(self):
|
||
"""Test basic command properties"""
|
||
command = DeleteCommand()
|
||
assert command.command_string == "delete"
|
||
assert "Delete data" in command.help_string
|
||
assert command.docs_url is not None
|
||
|
||
def test_configure_parser(self):
|
||
"""Test parser configuration"""
|
||
command = DeleteCommand()
|
||
parser = argparse.ArgumentParser()
|
||
|
||
command.configure_parser(parser)
|
||
|
||
# Check that arguments are added
|
||
actions = {action.dest: action for action in parser._actions}
|
||
assert "dataset_name" in actions
|
||
assert "user_id" in actions
|
||
assert "all" in actions
|
||
assert "force" in actions
|
||
|
||
@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):
|
||
"""Test execute delete dataset with user confirmation"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.delete = AsyncMock()
|
||
|
||
with patch.dict(sys.modules, {"cognee": mock_cognee}):
|
||
command = DeleteCommand()
|
||
args = argparse.Namespace(
|
||
dataset_name="test_dataset", user_id=None, all=False, force=False
|
||
)
|
||
|
||
mock_confirm.return_value = True
|
||
|
||
command.execute(args)
|
||
|
||
mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?")
|
||
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)
|
||
|
||
@patch("cognee.cli.commands.delete_command.fmt.confirm")
|
||
def test_execute_delete_cancelled(self, mock_confirm):
|
||
"""Test execute when user cancels deletion"""
|
||
command = DeleteCommand()
|
||
args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=False)
|
||
|
||
mock_confirm.return_value = False
|
||
|
||
# Should not raise exception, just return
|
||
command.execute(args)
|
||
|
||
mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?")
|
||
|
||
@patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
|
||
def test_execute_delete_forced(self, mock_asyncio_run):
|
||
"""Test execute delete with force flag"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.delete = AsyncMock()
|
||
|
||
with patch.dict(sys.modules, {"cognee": mock_cognee}):
|
||
command = DeleteCommand()
|
||
args = argparse.Namespace(
|
||
dataset_name="test_dataset", user_id=None, all=False, force=True
|
||
)
|
||
|
||
command.execute(args)
|
||
|
||
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)
|
||
|
||
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)
|
||
|
||
# Should not raise exception, just return with error message
|
||
command.execute(args)
|
||
|
||
@patch("cognee.cli.commands.delete_command.asyncio.run")
|
||
def test_execute_with_exception(self, mock_asyncio_run):
|
||
"""Test execute handles exceptions properly"""
|
||
command = DeleteCommand()
|
||
args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=True)
|
||
|
||
mock_asyncio_run.side_effect = Exception("Delete error")
|
||
|
||
with pytest.raises(CliCommandException):
|
||
command.execute(args)
|
||
|
||
|
||
class TestConfigCommand:
|
||
"""Test the ConfigCommand class"""
|
||
|
||
def test_command_properties(self):
|
||
"""Test basic command properties"""
|
||
command = ConfigCommand()
|
||
assert command.command_string == "config"
|
||
assert "Manage cognee configuration" in command.help_string
|
||
assert command.docs_url is not None
|
||
|
||
def test_configure_parser(self):
|
||
"""Test parser configuration"""
|
||
command = ConfigCommand()
|
||
parser = argparse.ArgumentParser()
|
||
|
||
command.configure_parser(parser)
|
||
|
||
# Check that subparsers are created
|
||
subparsers_actions = [
|
||
action for action in parser._actions if isinstance(action, argparse._SubParsersAction)
|
||
]
|
||
assert len(subparsers_actions) == 1
|
||
|
||
subparsers = subparsers_actions[0]
|
||
assert "get" in subparsers.choices
|
||
assert "set" in subparsers.choices
|
||
assert "list" in subparsers.choices
|
||
assert "unset" in subparsers.choices
|
||
assert "reset" in subparsers.choices
|
||
|
||
def test_execute_no_action(self):
|
||
"""Test execute when no config action is provided"""
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace()
|
||
|
||
# Should not raise exception, just return with error message
|
||
command.execute(args)
|
||
|
||
@patch("builtins.__import__")
|
||
def test_execute_get_action(self, mock_import):
|
||
"""Test execute get action"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.config.get = MagicMock(return_value="openai")
|
||
mock_import.return_value = mock_cognee
|
||
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace(config_action="get", key="llm_provider")
|
||
|
||
command.execute(args)
|
||
|
||
@patch("builtins.__import__")
|
||
def test_execute_set_action(self, mock_import):
|
||
"""Test execute set action"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.config.set = MagicMock()
|
||
mock_import.return_value = mock_cognee
|
||
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace(config_action="set", key="llm_provider", value="anthropic")
|
||
|
||
command.execute(args)
|
||
|
||
@patch("builtins.__import__")
|
||
def test_execute_set_action_json_value(self, mock_import):
|
||
"""Test execute set action with JSON value"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.config.set = MagicMock()
|
||
mock_import.return_value = mock_cognee
|
||
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace(config_action="set", key="chunk_size", value="1024")
|
||
|
||
command.execute(args)
|
||
|
||
def test_execute_list_action(self):
|
||
"""Test execute list action"""
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace(config_action="list")
|
||
|
||
# Should not raise exception
|
||
command.execute(args)
|
||
|
||
@patch("cognee.cli.commands.config_command.fmt.confirm")
|
||
def test_execute_unset_action(self, mock_confirm):
|
||
"""Test execute unset action"""
|
||
# Mock the cognee module
|
||
mock_cognee = MagicMock()
|
||
mock_cognee.config.set_llm_provider = MagicMock()
|
||
|
||
with patch.dict(sys.modules, {"cognee": mock_cognee}):
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace(config_action="unset", key="llm_provider", force=False)
|
||
|
||
mock_confirm.return_value = True
|
||
|
||
command.execute(args)
|
||
|
||
mock_confirm.assert_called_once()
|
||
|
||
@patch("cognee.cli.commands.config_command.fmt.confirm")
|
||
def test_execute_reset_action(self, mock_confirm):
|
||
"""Test execute reset action"""
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace(config_action="reset", force=False)
|
||
|
||
mock_confirm.return_value = True
|
||
|
||
# Should not raise exception
|
||
command.execute(args)
|
||
|
||
mock_confirm.assert_called_once()
|
||
|
||
def test_execute_with_exception(self):
|
||
"""Test execute handles exceptions properly"""
|
||
# Test with an invalid action that will cause an exception in the main execute method
|
||
command = ConfigCommand()
|
||
args = argparse.Namespace(config_action="invalid_action")
|
||
|
||
# This should not raise CliCommandException, just handle it gracefully
|
||
# The config command handles unknown actions by showing an error message
|
||
command.execute(args)
|