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

cognee/tests/integration/cli/test_cli_integration.py

@@ -0,0 +1,331 @@
+"""
+Integration tests for CLI commands that test end-to-end functionality.
+"""
+
+import tempfile
+import os
+import sys
+import subprocess
+from pathlib import Path
+from unittest.mock import patch, MagicMock
+
+
+class TestCliIntegration:
+    """Integration tests for CLI commands"""
+
+    def test_cli_help(self):
+        """Test that CLI help works"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "--help"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        assert result.returncode == 0
+        assert "cognee" in result.stdout.lower()
+        assert "available commands" in result.stdout.lower()
+
+    def test_cli_version(self):
+        """Test that CLI version works"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "--version"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        assert result.returncode == 0
+        assert "cognee" in result.stdout.lower()
+
+    def test_command_help(self):
+        """Test that individual command help works"""
+        commands = ["add", "search", "cognify", "delete", "config"]
+
+        for command in commands:
+            result = subprocess.run(
+                [sys.executable, "-m", "cognee.cli._cognee", command, "--help"],
+                capture_output=True,
+                text=True,
+                cwd=Path(__file__).parent.parent.parent,  # Go to project root
+            )
+
+            assert result.returncode == 0, f"Command {command} help failed"
+            assert command in result.stdout.lower()
+
+    def test_invalid_command(self):
+        """Test that invalid commands are handled properly"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "invalid_command"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        assert result.returncode != 0
+
+    @patch("cognee.add")
+    def test_add_command_integration(self, mock_add):
+        """Test add command integration"""
+        mock_add.return_value = None
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
+            f.write("Test content for CLI integration")
+            temp_file = f.name
+
+        try:
+            result = subprocess.run(
+                [sys.executable, "-m", "cognee.cli._cognee", "add", temp_file],
+                capture_output=True,
+                text=True,
+                cwd=Path(__file__).parent.parent.parent,  # Go to project root
+            )
+
+            # Note: This might fail due to dependencies, but we're testing the CLI structure
+            # The important thing is that it doesn't crash with argument parsing errors
+            # Allow litellm logging worker cancellation errors as they're expected during process shutdown
+            stderr_lower = result.stderr.lower()
+            has_error = "error" in stderr_lower
+            has_expected_failure = "failed to add data" in stderr_lower
+            has_litellm_cancellation = (
+                "loggingworker cancelled" in stderr_lower or "cancellederror" in stderr_lower
+            )
+
+            assert not has_error or has_expected_failure or has_litellm_cancellation
+
+        finally:
+            os.unlink(temp_file)
+
+    def test_config_subcommands(self):
+        """Test config subcommands help"""
+        subcommands = ["get", "set", "list", "unset", "reset"]
+
+        for subcommand in subcommands:
+            result = subprocess.run(
+                [sys.executable, "-m", "cognee.cli._cognee", "config", subcommand, "--help"],
+                capture_output=True,
+                text=True,
+                cwd=Path(__file__).parent.parent.parent,  # Go to project root
+            )
+
+            assert result.returncode == 0, f"Config {subcommand} help failed"
+
+    def test_search_command_missing_query(self):
+        """Test search command fails when query is missing"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "search"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        assert result.returncode != 0
+        assert "required" in result.stderr.lower() or "error" in result.stderr.lower()
+
+    def test_delete_command_no_target(self):
+        """Test delete command with no target specified"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "delete"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        # Should run but show error message about missing target
+        # Return code might be 0 since the command handles this gracefully
+        assert (
+            "specify what to delete" in result.stdout.lower()
+            or "specify what to delete" in result.stderr.lower()
+        )
+
+
+class TestCliArgumentParsing:
+    """Test CLI argument parsing edge cases"""
+
+    def test_add_multiple_files(self):
+        """Test add command with multiple file arguments"""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file1 = os.path.join(temp_dir, "file1.txt")
+            file2 = os.path.join(temp_dir, "file2.txt")
+
+            with open(file1, "w") as f:
+                f.write("Content 1")
+            with open(file2, "w") as f:
+                f.write("Content 2")
+
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "cognee.cli._cognee",
+                    "add",
+                    file1,
+                    file2,
+                    "--dataset-name",
+                    "test",
+                ],
+                capture_output=True,
+                text=True,
+                cwd=Path(__file__).parent.parent.parent,  # Go to project root
+            )
+
+            # Test that argument parsing works (regardless of actual execution)
+            assert (
+                "argument" not in result.stderr.lower() or "failed to add" in result.stderr.lower()
+            )
+
+    def test_search_with_all_options(self):
+        """Test search command with all possible options"""
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "cognee.cli._cognee",
+                "search",
+                "test query",
+                "--query-type",
+                "CHUNKS",
+                "--datasets",
+                "dataset1",
+                "dataset2",
+                "--top-k",
+                "5",
+                "--output-format",
+                "json",
+            ],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        # Should not have argument parsing errors
+        assert "unrecognized arguments" not in result.stderr.lower()
+        assert "invalid choice" not in result.stderr.lower()
+
+    def test_cognify_with_all_options(self):
+        """Test cognify command with all possible options"""
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "cognee.cli._cognee",
+                "cognify",
+                "--datasets",
+                "dataset1",
+                "dataset2",
+                "--chunk-size",
+                "1024",
+                "--chunker",
+                "TextChunker",
+                "--background",
+                "--verbose",
+            ],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        # Should not have argument parsing errors
+        assert "unrecognized arguments" not in result.stderr.lower()
+        assert "invalid choice" not in result.stderr.lower()
+
+    def test_config_set_command(self):
+        """Test config set command argument parsing"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "config", "set", "test_key", "test_value"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        # Should not have argument parsing errors
+        assert "unrecognized arguments" not in result.stderr.lower()
+        assert "required" not in result.stderr.lower() or "failed to set" in result.stderr.lower()
+
+    def test_delete_with_force(self):
+        """Test delete command with force flag"""
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "cognee.cli._cognee",
+                "delete",
+                "--dataset-name",
+                "test_dataset",
+                "--force",
+            ],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        # Should not have argument parsing errors
+        assert "unrecognized arguments" not in result.stderr.lower()
+
+
+class TestCliErrorHandling:
+    """Test CLI error handling and edge cases"""
+
+    def test_debug_mode_flag(self):
+        """Test that debug flag is accepted"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "--debug", "search", "test query"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        # Should not have argument parsing errors for debug flag
+        assert "unrecognized arguments" not in result.stderr.lower()
+
+    def test_invalid_search_type(self):
+        """Test invalid search type handling"""
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "cognee.cli._cognee",
+                "search",
+                "test query",
+                "--query-type",
+                "INVALID_TYPE",
+            ],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        assert result.returncode != 0
+        assert "invalid choice" in result.stderr.lower()
+
+    def test_invalid_chunker(self):
+        """Test invalid chunker handling"""
+        result = subprocess.run(
+            [sys.executable, "-m", "cognee.cli._cognee", "cognify", "--chunker", "InvalidChunker"],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        assert result.returncode != 0
+        assert "invalid choice" in result.stderr.lower()
+
+    def test_invalid_output_format(self):
+        """Test invalid output format handling"""
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "cognee.cli._cognee",
+                "search",
+                "test query",
+                "--output-format",
+                "invalid",
+            ],
+            capture_output=True,
+            text=True,
+            cwd=Path(__file__).parent.parent.parent,  # Go to project root
+        )
+
+        assert result.returncode != 0
+        assert "invalid choice" in result.stderr.lower()

cognee/tests/integration/documents/PdfDocument_test.py

@@ -2,6 +2,7 @@ import os
 import sys
 import uuid
 import pytest
+import pathlib
 from unittest.mock import patch
 
 from cognee.modules.chunking.TextChunker import TextChunker
@@ -24,8 +25,7 @@ GROUND_TRUTH = [
 @pytest.mark.asyncio
 async def test_PdfDocument(mock_engine):
     test_file_path = os.path.join(
-        os.sep,
-        *(os.path.dirname(__file__).split(os.sep)[:-2]),
+        pathlib.Path(__file__).parent.parent.parent,
         "test_data",
         "artificial-intelligence.pdf",
     )

cognee/tests/integration/documents/TextDocument_test.py

@@ -2,6 +2,7 @@ import os
 import sys
 import uuid
 import pytest
+import pathlib
 from unittest.mock import patch
 
 from cognee.modules.chunking.TextChunker import TextChunker
@@ -34,10 +35,7 @@ GROUND_TRUTH = {
 @pytest.mark.asyncio
 async def test_TextDocument(mock_engine, input_file, chunk_size):
     test_file_path = os.path.join(
-        os.sep,
-        *(os.path.dirname(__file__).split(os.sep)[:-2]),
-        "test_data",
-        input_file,
+        pathlib.Path(__file__).parent.parent.parent, "test_data", input_file
     )
     document = TextDocument(
         id=uuid.uuid4(),

cognee/tests/integration/documents/UnstructuredDocument_test.py

@@ -2,6 +2,7 @@ import os
 import sys
 import uuid
 import pytest
+import pathlib
 from unittest.mock import patch
 
 from cognee.modules.chunking.TextChunker import TextChunker
@@ -18,29 +19,25 @@ chunk_by_sentence_module = sys.modules.get("cognee.tasks.chunks.chunk_by_sentenc
 async def test_UnstructuredDocument(mock_engine):
     # Define file paths of test data
     pptx_file_path = os.path.join(
-        os.sep,
-        *(os.path.dirname(__file__).split(os.sep)[:-2]),
+        pathlib.Path(__file__).parent.parent.parent,
         "test_data",
         "example.pptx",
     )
 
     docx_file_path = os.path.join(
-        os.sep,
-        *(os.path.dirname(__file__).split(os.sep)[:-2]),
+        pathlib.Path(__file__).parent.parent.parent,
         "test_data",
         "example.docx",
     )
 
     csv_file_path = os.path.join(
-        os.sep,
-        *(os.path.dirname(__file__).split(os.sep)[:-2]),
+        pathlib.Path(__file__).parent.parent.parent,
         "test_data",
         "example.csv",
     )
 
     xlsx_file_path = os.path.join(
-        os.sep,
-        *(os.path.dirname(__file__).split(os.sep)[:-2]),
+        pathlib.Path(__file__).parent.parent.parent,
         "test_data",
         "example.xlsx",
     )

cognee/tests/{test_deletion.py → test_delete_hard.py}

@@ -45,8 +45,6 @@ async def main():
     Each of these car manufacturer contributes to Germany's reputation as a leader in the global automotive industry, showcasing a blend of innovation, performance, and design excellence.
     """
 
-    ################### HARD DELETE
-
     # Add documents and get dataset information
     add_result = await cognee.add(
         [
@@ -80,41 +78,6 @@ async def main():
 
     assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with hard delete."
 
-    ################### SOFT DELETE
-
-    # Add documents and get dataset information
-    add_result = await cognee.add(
-        [
-            pdf_document,
-            txt_document,
-            text_document_as_literal,
-            unstructured_document,
-            audio_document,
-            image_document,
-        ]
-    )
-    dataset_id = add_result.dataset_id
-
-    await cognee.cognify()
-
-    from cognee.infrastructure.databases.graph import get_graph_engine
-
-    graph_engine = await get_graph_engine()
-    nodes, edges = await graph_engine.get_graph_data()
-    assert len(nodes) > 10 and len(edges) > 10, "Graph database is not loaded."
-
-    # Get the data IDs from the dataset
-    dataset_data = await get_dataset_data(dataset_id)
-    assert len(dataset_data) > 0, "Dataset should contain data"
-
-    # Delete each document using its ID
-    for data_item in dataset_data:
-        await cognee.delete(data_item.id, dataset_id, mode="soft")
-
-    nodes, edges = await graph_engine.get_graph_data()
-
-    assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with soft delete."
-
 
 if __name__ == "__main__":
     import asyncio

cognee/tests/test_delete_soft.py

@@ -0,0 +1,85 @@
+import os
+import shutil
+import cognee
+import pathlib
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.data.methods import get_dataset_data
+
+logger = get_logger()
+
+
+async def main():
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    pdf_document = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/artificial-intelligence.pdf"
+    )
+
+    txt_document = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Natural_language_processing_copy.txt"
+    )
+
+    audio_document = os.path.join(pathlib.Path(__file__).parent, "test_data/text_to_speech.mp3")
+
+    image_document = os.path.join(pathlib.Path(__file__).parent, "test_data/example.png")
+
+    unstructured_document = os.path.join(pathlib.Path(__file__).parent, "test_data/example.pptx")
+
+    text_document_as_literal = """
+    1. Audi
+    Audi is known for its modern designs and advanced technology. Founded in the early 1900s, the brand has earned a reputation for precision engineering and innovation. With features like the Quattro all-wheel-drive system, Audi offers a range of vehicles from stylish sedans to high-performance sports cars.
+
+    2. BMW
+    BMW, short for Bayerische Motoren Werke, is celebrated for its focus on performance and driving pleasure. The company's vehicles are designed to provide a dynamic and engaging driving experience, and their slogan, "The Ultimate Driving Machine," reflects that commitment. BMW produces a variety of cars that combine luxury with sporty performance.
+
+    3. Mercedes-Benz
+    Mercedes-Benz is synonymous with luxury and quality. With a history dating back to the early 20th century, the brand is known for its elegant designs, innovative safety features, and high-quality engineering. Mercedes-Benz manufactures not only luxury sedans but also SUVs, sports cars, and commercial vehicles, catering to a wide range of needs.
+
+    4. Porsche
+    Porsche is a name that stands for high-performance sports cars. Founded in 1931, the brand has become famous for models like the iconic Porsche 911. Porsche cars are celebrated for their speed, precision, and distinctive design, appealing to car enthusiasts who value both performance and style.
+
+    5. Volkswagen
+    Volkswagen, which means "people's car" in German, was established with the idea of making affordable and reliable vehicles accessible to everyone. Over the years, Volkswagen has produced several iconic models, such as the Beetle and the Golf. Today, it remains one of the largest car manufacturers in the world, offering a wide range of vehicles that balance practicality with quality.
+
+    Each of these car manufacturer contributes to Germany's reputation as a leader in the global automotive industry, showcasing a blend of innovation, performance, and design excellence.
+    """
+
+    # Add documents and get dataset information
+    add_result = await cognee.add(
+        [
+            pdf_document,
+            txt_document,
+            text_document_as_literal,
+            unstructured_document,
+            audio_document,
+            image_document,
+        ]
+    )
+    dataset_id = add_result.dataset_id
+
+    await cognee.cognify()
+
+    from cognee.infrastructure.databases.graph import get_graph_engine
+
+    graph_engine = await get_graph_engine()
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) > 10 and len(edges) > 10, "Graph database is not loaded."
+
+    # Get the data IDs from the dataset
+    dataset_data = await get_dataset_data(dataset_id)
+    assert len(dataset_data) > 0, "Dataset should contain data"
+
+    # Delete each document using its ID
+    for data_item in dataset_data:
+        await cognee.delete(data_item.id, dataset_id, mode="soft")
+
+    nodes, edges = await graph_engine.get_graph_data()
+
+    assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with soft delete."
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

cognee/tests/test_kuzu.py

@@ -104,11 +104,11 @@ async def main():
             node_name=["nonexistent"],
         ).get_context("What is in the context?")
 
-        assert isinstance(context_nonempty, str) and context_nonempty != "", (
+        assert isinstance(context_nonempty, list) and context_nonempty != [], (
             f"Nodeset_search_test:Expected non-empty string for context_nonempty, got: {context_nonempty!r}"
         )
 
-        assert context_empty == "", (
+        assert context_empty == [], (
             f"Nodeset_search_test:Expected empty string for context_empty, got: {context_empty!r}"
         )
 

cognee/tests/test_neo4j.py

@@ -108,11 +108,11 @@ async def main():
         node_name=["nonexistent"],
     ).get_context("What is in the context?")
 
-    assert isinstance(context_nonempty, str) and context_nonempty != "", (
+    assert isinstance(context_nonempty, list) and context_nonempty != [], (
         f"Nodeset_search_test:Expected non-empty string for context_nonempty, got: {context_nonempty!r}"
     )
 
-    assert context_empty == "", (
+    assert context_empty == [], (
         f"Nodeset_search_test:Expected empty string for context_empty, got: {context_empty!r}"
     )
 

cognee/tests/test_permissions.py

@@ -79,7 +79,7 @@ async def main():
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
-    assert search_results[0]["dataset_name"] == "NLP", (
+    assert search_results[0].dataset_name == "NLP", (
         f"Dict must contain dataset name 'NLP': {search_results[0]}"
     )
 
@@ -93,7 +93,7 @@ async def main():
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
-    assert search_results[0]["dataset_name"] == "QUANTUM", (
+    assert search_results[0].dataset_name == "QUANTUM", (
         f"Dict must contain dataset name 'QUANTUM': {search_results[0]}"
     )
 
@@ -170,7 +170,7 @@ async def main():
     for result in search_results:
         print(f"{result}\n")
 
-    assert search_results[0]["dataset_name"] == "QUANTUM", (
+    assert search_results[0].dataset_name == "QUANTUM", (
         f"Dict must contain dataset name 'QUANTUM': {search_results[0]}"
     )
 

cognee/tests/test_relational_db_migration.py

@@ -1,6 +1,6 @@
-import json
 import pathlib
 import os
+from typing import List
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.relational import (
     get_migration_relational_engine,
@@ -10,7 +10,7 @@ from cognee.infrastructure.databases.vector.pgvector import (
     create_db_and_tables as create_pgvector_db_and_tables,
 )
 from cognee.tasks.ingestion import migrate_relational_database
-from cognee.modules.search.types import SearchType
+from cognee.modules.search.types import SearchResult, SearchType
 import cognee
 
 
@@ -45,13 +45,15 @@ async def relational_db_migration():
     await migrate_relational_database(graph_engine, schema=schema)
 
     # 1. Search the graph
-    search_results = await cognee.search(
+    search_results: List[SearchResult] = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION, query_text="Tell me about the artist AC/DC"
-    )
+    )  # type: ignore
     print("Search results:", search_results)
 
     # 2. Assert that the search results contain "AC/DC"
-    assert any("AC/DC" in r for r in search_results), "AC/DC not found in search results!"
+    assert any("AC/DC" in r.search_result for r in search_results), (
+        "AC/DC not found in search results!"
+    )
 
     migration_db_provider = migration_engine.engine.dialect.name
     if migration_db_provider == "postgresql":