cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0__py3-none-any.whl

cognee/tests/unit/cli/test_cli_main.py

@@ -0,0 +1,173 @@
+"""
+Tests for the main CLI entry point and command discovery.
+"""
+
+import pytest
+import argparse
+from unittest.mock import patch, MagicMock
+from cognee.cli._cognee import main, _discover_commands, _create_parser
+from cognee.cli.exceptions import CliCommandException, CliCommandInnerException
+
+
+class TestCliMain:
+    """Test the main CLI functionality"""
+
+    def test_discover_commands(self):
+        """Test that all expected commands are discovered"""
+        commands = _discover_commands()
+
+        # Check that we get command classes back
+        assert len(commands) > 0
+
+        # Check that we have the expected commands
+        command_strings = []
+        for command_class in commands:
+            command = command_class()
+            command_strings.append(command.command_string)
+
+        expected_commands = ["add", "search", "cognify", "delete", "config"]
+        for expected_command in expected_commands:
+            assert expected_command in command_strings
+
+    def test_create_parser(self):
+        """Test parser creation and command installation"""
+        parser, installed_commands = _create_parser()
+
+        # Check parser is created
+        assert isinstance(parser, argparse.ArgumentParser)
+
+        # Check commands are installed
+        expected_commands = ["add", "search", "cognify", "delete", "config"]
+        for expected_command in expected_commands:
+            assert expected_command in installed_commands
+
+        # Check parser has version argument
+        actions = [action.dest for action in parser._actions]
+        assert "version" in actions
+
+    @patch("cognee.cli._cognee._create_parser")
+    def test_main_no_command(self, mock_create_parser):
+        """Test main function when no command is provided"""
+        mock_parser = MagicMock()
+        mock_parser.parse_args.return_value = MagicMock(command=None)
+        mock_create_parser.return_value = (mock_parser, {})
+
+        result = main()
+
+        assert result == -1
+        mock_parser.print_help.assert_called_once()
+
+    @patch("cognee.cli._cognee._create_parser")
+    def test_main_with_valid_command(self, mock_create_parser):
+        """Test main function with a valid command"""
+        mock_command = MagicMock()
+        mock_command.execute.return_value = None
+
+        mock_parser = MagicMock()
+        mock_args = MagicMock(command="test")
+        mock_parser.parse_args.return_value = mock_args
+
+        mock_create_parser.return_value = (mock_parser, {"test": mock_command})
+
+        result = main()
+
+        assert result == 0
+        mock_command.execute.assert_called_once_with(mock_args)
+
+    @patch("cognee.cli._cognee._create_parser")
+    @patch("cognee.cli.debug.is_debug_enabled")
+    def test_main_with_command_exception(self, mock_debug, mock_create_parser):
+        """Test main function when command raises exception"""
+        mock_debug.return_value = False
+
+        mock_command = MagicMock()
+        mock_command.execute.side_effect = CliCommandException("Test error", error_code=2)
+
+        mock_parser = MagicMock()
+        mock_args = MagicMock(command="test")
+        mock_parser.parse_args.return_value = mock_args
+
+        mock_create_parser.return_value = (mock_parser, {"test": mock_command})
+
+        result = main()
+
+        assert result == 2
+
+    @patch("cognee.cli._cognee._create_parser")
+    @patch("cognee.cli.debug.is_debug_enabled")
+    def test_main_with_generic_exception(self, mock_debug, mock_create_parser):
+        """Test main function when command raises generic exception"""
+        mock_debug.return_value = False
+
+        mock_command = MagicMock()
+        mock_command.execute.side_effect = Exception("Generic error")
+
+        mock_parser = MagicMock()
+        mock_args = MagicMock(command="test")
+        mock_parser.parse_args.return_value = mock_args
+
+        mock_create_parser.return_value = (mock_parser, {"test": mock_command})
+
+        result = main()
+
+        assert result == -1
+
+    @patch("cognee.cli._cognee._create_parser")
+    @patch("cognee.cli.debug.is_debug_enabled")
+    def test_main_debug_mode_reraises_exception(self, mock_debug, mock_create_parser):
+        """Test main function reraises exceptions in debug mode"""
+        mock_debug.return_value = True
+
+        test_exception = CliCommandException(
+            "Test error", error_code=2, raiseable_exception=ValueError("Inner error")
+        )
+
+        mock_command = MagicMock()
+        mock_command.execute.side_effect = test_exception
+
+        mock_parser = MagicMock()
+        mock_args = MagicMock(command="test")
+        mock_parser.parse_args.return_value = mock_args
+
+        mock_create_parser.return_value = (mock_parser, {"test": mock_command})
+
+        with pytest.raises(ValueError, match="Inner error"):
+            main()
+
+    def test_version_argument(self):
+        """Test that version argument is properly configured"""
+        parser, _ = _create_parser()
+
+        # Check that version action exists
+        version_actions = [action for action in parser._actions if action.dest == "version"]
+        assert len(version_actions) == 1
+
+        version_action = version_actions[0]
+        assert "cognee" in version_action.version
+
+    def test_debug_argument(self):
+        """Test that debug argument is properly configured"""
+        parser, _ = _create_parser()
+
+        # Check that debug action exists
+        debug_actions = [action for action in parser._actions if action.dest == "debug"]
+        assert len(debug_actions) == 1
+
+
+class TestDebugAction:
+    """Test the DebugAction class"""
+
+    @patch("cognee.cli.debug.enable_debug")
+    @patch("cognee.cli.echo.note")
+    def test_debug_action_call(self, mock_note, mock_enable_debug):
+        """Test that DebugAction enables debug mode"""
+        from cognee.cli._cognee import DebugAction
+
+        action = DebugAction([])
+        parser = MagicMock()
+        namespace = MagicMock()
+
+        action(parser, namespace, None)
+
+        mock_enable_debug.assert_called_once()
+        mock_note.assert_called_once_with("Debug mode enabled. Full stack traces will be shown.")

cognee/tests/unit/cli/test_cli_runner.py

@@ -0,0 +1,62 @@
+"""
+Test runner and utilities for CLI tests.
+"""
+
+import pytest
+import sys
+from pathlib import Path
+
+
+def run_cli_tests():
+    """Run all CLI tests"""
+    test_dir = Path(__file__).parent
+    integration_dir = test_dir.parent.parent / "integration" / "cli"
+
+    cli_unit_test_files = [
+        "test_cli_main.py",
+        "test_cli_commands.py",
+        "test_cli_utils.py",
+        "test_cli_edge_cases.py",
+    ]
+
+    cli_integration_test_files = ["test_cli_integration.py"]
+
+    # Run tests with pytest
+    args = ["-v", "--tb=short"]
+
+    # Add unit tests
+    for test_file in cli_unit_test_files:
+        test_path = test_dir / test_file
+        if test_path.exists():
+            args.append(str(test_path))
+
+    # Add integration tests
+    for test_file in cli_integration_test_files:
+        test_path = integration_dir / test_file
+        if test_path.exists():
+            args.append(str(test_path))
+
+    return pytest.main(args)
+
+
+def run_specific_cli_test(test_file):
+    """Run a specific CLI test file"""
+    test_dir = Path(__file__).parent
+    test_path = test_dir / test_file
+
+    if not test_path.exists():
+        print(f"Test file {test_file} not found")
+        return 1
+
+    return pytest.main(["-v", "--tb=short", str(test_path)])
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1:
+        # Run specific test file
+        exit_code = run_specific_cli_test(sys.argv[1])
+    else:
+        # Run all CLI tests
+        exit_code = run_cli_tests()
+
+    sys.exit(exit_code)

cognee/tests/unit/cli/test_cli_utils.py

@@ -0,0 +1,127 @@
+"""
+Tests for CLI utility functions and helper modules.
+"""
+
+from cognee.cli import debug
+from cognee.cli.config import (
+    CLI_DESCRIPTION,
+    DEFAULT_DOCS_URL,
+    COMMAND_DESCRIPTIONS,
+    SEARCH_TYPE_CHOICES,
+    CHUNKER_CHOICES,
+    OUTPUT_FORMAT_CHOICES,
+)
+from cognee.cli._cognee import _discover_commands
+
+
+class TestCliConfig:
+    """Test CLI configuration constants"""
+
+    def test_cli_description_exists(self):
+        """Test that CLI description is defined"""
+        assert CLI_DESCRIPTION
+        assert isinstance(CLI_DESCRIPTION, str)
+        assert "cognee" in CLI_DESCRIPTION.lower()
+
+    def test_default_docs_url_exists(self):
+        """Test that default docs URL is defined"""
+        assert DEFAULT_DOCS_URL
+        assert isinstance(DEFAULT_DOCS_URL, str)
+        assert DEFAULT_DOCS_URL.startswith("https://")
+        assert "cognee.ai" in DEFAULT_DOCS_URL
+
+    def test_command_descriptions_complete(self):
+        """Test that all expected commands have descriptions"""
+        commands = _discover_commands()
+        assert len(commands) > 0
+
+        expected_commands = []
+        for command_class in commands:
+            command = command_class()
+            expected_commands.append(command.command_string)
+
+        for command in expected_commands:
+            assert command in COMMAND_DESCRIPTIONS
+            assert isinstance(COMMAND_DESCRIPTIONS[command], str)
+            assert len(COMMAND_DESCRIPTIONS[command]) > 0
+
+    def test_search_type_choices_valid(self):
+        """Test that search type choices are valid"""
+        assert isinstance(SEARCH_TYPE_CHOICES, list)
+        assert len(SEARCH_TYPE_CHOICES) > 0
+
+        expected_types = [
+            "GRAPH_COMPLETION",
+            "RAG_COMPLETION",
+            "INSIGHTS",
+            "CHUNKS",
+            "SUMMARIES",
+            "CODE",
+            "CYPHER",
+        ]
+
+        for expected_type in expected_types:
+            assert expected_type in SEARCH_TYPE_CHOICES
+
+    def test_chunker_choices_valid(self):
+        """Test that chunker choices are valid"""
+        assert isinstance(CHUNKER_CHOICES, list)
+        assert len(CHUNKER_CHOICES) > 0
+        assert "TextChunker" in CHUNKER_CHOICES
+        assert "LangchainChunker" in CHUNKER_CHOICES
+
+    def test_output_format_choices_valid(self):
+        """Test that output format choices are valid"""
+        assert isinstance(OUTPUT_FORMAT_CHOICES, list)
+        assert len(OUTPUT_FORMAT_CHOICES) > 0
+
+        expected_formats = ["json", "pretty", "simple"]
+        for expected_format in expected_formats:
+            assert expected_format in OUTPUT_FORMAT_CHOICES
+
+
+class TestCliReference:
+    """Test CLI reference protocol"""
+
+    def test_supports_cli_command_protocol(self):
+        """Test that SupportsCliCommand protocol is properly defined"""
+        from cognee.cli.reference import SupportsCliCommand
+
+        # Test that it's a protocol
+        assert hasattr(SupportsCliCommand, "__annotations__")
+
+        # Test required attributes
+        annotations = SupportsCliCommand.__annotations__
+        assert "command_string" in annotations
+        assert "help_string" in annotations
+        assert "description" in annotations
+        assert "docs_url" in annotations
+
+    def test_protocol_methods(self):
+        """Test that protocol defines required methods"""
+        from cognee.cli.reference import SupportsCliCommand
+        import inspect
+
+        # Get abstract methods
+        abstract_methods = []
+        for name, method in inspect.getmembers(SupportsCliCommand, predicate=inspect.ismethod):
+            if getattr(method, "__isabstractmethod__", False):
+                abstract_methods.append(name)
+
+        # Should have abstract methods for configure_parser and execute
+        method_names = [name for name, _ in inspect.getmembers(SupportsCliCommand)]
+        assert "configure_parser" in method_names
+        assert "execute" in method_names
+
+
+class TestCliUtilityFunctions:
+    """Test utility functions and edge cases"""
+
+    def test_multiple_debug_enable_calls(self):
+        """Test multiple calls to enable_debug"""
+        debug.enable_debug()
+        debug.enable_debug()  # Should not cause issues
+        assert debug.is_debug_enabled() is True
+
+        # Reset for other tests
+        debug._debug_enabled = False

cognee/tests/unit/modules/retrieval/chunks_retriever_test.py

@@ -1,7 +1,7 @@
 import os
 import pytest
 import pathlib
-
+from typing import List
 import cognee
 from cognee.low_level import setup
 from cognee.tasks.storage import add_data_points
@@ -10,6 +10,20 @@ from cognee.modules.chunking.models import DocumentChunk
 from cognee.modules.data.processing.document_types import TextDocument
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.data.processing.document_types import Document
+from cognee.modules.engine.models import Entity
+
+
+class DocumentChunkWithEntities(DataPoint):
+    text: str
+    chunk_size: int
+    chunk_index: int
+    cut_type: str
+    is_part_of: Document
+    contains: List[Entity] = None
+
+    metadata: dict = {"index_fields": ["text"]}
 
 
 class TestChunksRetriever:
@@ -179,7 +193,9 @@ class TestChunksRetriever:
             await retriever.get_context("Christina Mayer")
 
         vector_engine = get_vector_engine()
-        await vector_engine.create_collection("DocumentChunk_text", payload_schema=DocumentChunk)
+        await vector_engine.create_collection(
+            "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
+        )
 
         context = await retriever.get_context("Christina Mayer")
         assert len(context) == 0, "Found chunks when none should exist"

cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py

@@ -6,6 +6,7 @@ from typing import Optional, Union
 import cognee
 from cognee.low_level import setup, DataPoint
 from cognee.tasks.storage import add_data_points
+from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
     GraphCompletionContextExtensionRetriever,
@@ -51,17 +52,15 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
 
         answer = await retriever.get_completion("Who works at Canva?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_graph_completion_extension_context_complex(self):
@@ -129,7 +128,9 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context = await resolve_edges_to_text(
+            await retriever.get_context("Who works at Figma and drives Tesla?")
+        )
 
         print(context)
 
@@ -139,10 +140,8 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_get_graph_completion_extension_context_on_empty_graph(self):
@@ -168,11 +167,9 @@ class TestGraphCompletionWithContextExtensionRetriever:
         await setup()
 
         context = await retriever.get_context("Who works at Figma?")
-        assert context == "", "Context should be empty on an empty graph"
+        assert context == [], "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"

cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py

@@ -5,6 +5,7 @@ from typing import Optional, Union
 
 import cognee
 from cognee.low_level import setup, DataPoint
+from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
 from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
@@ -47,17 +48,15 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
 
         assert "Mike Broski --[works_for]--> Canva" in context, "Failed to get Mike Broski"
         assert "Christina Mayer --[works_for]--> Canva" in context, "Failed to get Christina Mayer"
 
         answer = await retriever.get_completion("Who works at Canva?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_graph_completion_cot_context_complex(self):
@@ -124,7 +123,7 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Figma?"))
 
         print(context)
 
@@ -134,10 +133,8 @@ class TestGraphCompletionCoTRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"
 
     @pytest.mark.asyncio
     async def test_get_graph_completion_cot_context_on_empty_graph(self):
@@ -163,11 +160,9 @@ class TestGraphCompletionCoTRetriever:
         await setup()
 
         context = await retriever.get_context("Who works at Figma?")
-        assert context == "", "Context should be empty on an empty graph"
+        assert context == [], "Context should be empty on an empty graph"
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
-        assert all(isinstance(item, str) and item.strip() for item in answer), (
-            "Answer must contain only non-empty strings"
-        )
+        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
+        assert answer.strip(), "Answer must contain only non-empty strings"

cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py

@@ -5,6 +5,7 @@ from typing import Optional, Union
 
 import cognee
 from cognee.low_level import setup, DataPoint
+from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
 from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
@@ -67,7 +68,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever()
 
-        context = await retriever.get_context("Who works at Canva?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Canva?"))
 
         # Ensure the top-level sections are present
         assert "Nodes:" in context, "Missing 'Nodes:' section in context"
@@ -191,7 +192,7 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever(top_k=20)
 
-        context = await retriever.get_context("Who works at Figma?")
+        context = await resolve_edges_to_text(await retriever.get_context("Who works at Figma?"))
 
         print(context)
 
@@ -223,4 +224,4 @@ class TestGraphCompletionRetriever:
         await setup()
 
         context = await retriever.get_context("Who works at Figma?")
-        assert context == "", "Context should be empty on an empty graph"
+        assert context == [], "Context should be empty on an empty graph"

cognee/tests/unit/modules/retrieval/insights_retriever_test.py

@@ -82,7 +82,7 @@ class TestInsightsRetriever:
 
         context = await retriever.get_context("Mike")
 
-        assert context[0][0]["name"] == "Mike Broski", "Failed to get Mike Broski"
+        assert context[0].node1.attributes["name"] == "Mike Broski", "Failed to get Mike Broski"
 
     @pytest.mark.asyncio
     async def test_insights_context_complex(self):
@@ -222,7 +222,9 @@ class TestInsightsRetriever:
 
         context = await retriever.get_context("Christina")
 
-        assert context[0][0]["name"] == "Christina Mayer", "Failed to get Christina Mayer"
+        assert context[0].node1.attributes["name"] == "Christina Mayer", (
+            "Failed to get Christina Mayer"
+        )
 
     @pytest.mark.asyncio
     async def test_insights_context_on_empty_graph(self):

cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py

@@ -1,7 +1,7 @@
 import os
+from typing import List
 import pytest
 import pathlib
-
 import cognee
 from cognee.low_level import setup
 from cognee.tasks.storage import add_data_points
@@ -10,6 +10,20 @@ from cognee.modules.chunking.models import DocumentChunk
 from cognee.modules.data.processing.document_types import TextDocument
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.modules.retrieval.completion_retriever import CompletionRetriever
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.data.processing.document_types import Document
+from cognee.modules.engine.models import Entity
+
+
+class DocumentChunkWithEntities(DataPoint):
+    text: str
+    chunk_size: int
+    chunk_index: int
+    cut_type: str
+    is_part_of: Document
+    contains: List[Entity] = None
+
+    metadata: dict = {"index_fields": ["text"]}
 
 
 class TestRAGCompletionRetriever:
@@ -182,7 +196,9 @@ class TestRAGCompletionRetriever:
             await retriever.get_context("Christina Mayer")
 
         vector_engine = get_vector_engine()
-        await vector_engine.create_collection("DocumentChunk_text", payload_schema=DocumentChunk)
+        await vector_engine.create_collection(
+            "DocumentChunk_text", payload_schema=DocumentChunkWithEntities
+        )
 
         context = await retriever.get_context("Christina Mayer")
         assert context == "", "Returned context should be empty on an empty graph"