Fix tests failing for python 3.10

cognee/tests/unit/cli/test_cli_commands.py

@@ -3,9 +3,10 @@ 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
+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
@@ -14,6 +15,16 @@ 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"""
 
@@ -39,58 +50,39 @@ class TestAddCommand:
         # Check data argument accepts multiple values
         assert actions["data"].nargs == "+"
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.add_command.asyncio")
+    def test_execute_single_item(self, mock_asyncio_run):
-    def test_execute_single_item(self, mock_asyncio, mock_import):
         """Test execute with single data item"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.add = AsyncMock()
-        mock_import.return_value = mock_cognee
 
-        # Mock asyncio.run to properly handle coroutines
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        def mock_run(coro):
+            command = AddCommand()
-            # Create an event loop and run the coroutine
+            args = argparse.Namespace(data=["test.txt"], dataset_name="test_dataset")
-            loop = asyncio.new_event_loop()
+            command.execute(args)
-            try:
-                return loop.run_until_complete(coro)
-            finally:
-                loop.close()
 
-        mock_asyncio.run.side_effect = mock_run
+        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")
 
-        command = AddCommand()
+    @patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run)
-        args = argparse.Namespace(data=["test.txt"], dataset_name="test_dataset")
+    def test_execute_multiple_items(self, mock_asyncio_run):
-
-        command.execute(args)
-
-        mock_asyncio.run.assert_called_once()
-
-    @patch("builtins.__import__")
-    @patch("cognee.cli.commands.add_command.asyncio")
-    def test_execute_multiple_items(self, mock_asyncio, mock_import):
         """Test execute with multiple data items"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.add = AsyncMock()
-        mock_import.return_value = mock_cognee
 
-        # Mock asyncio.run to properly handle coroutines
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        def mock_run(coro):
+            command = AddCommand()
-            loop = asyncio.new_event_loop()
+            args = argparse.Namespace(data=["test1.txt", "test2.txt"], dataset_name="test_dataset")
-            try:
+            command.execute(args)
-                return loop.run_until_complete(coro)
-            finally:
-                loop.close()
 
-        mock_asyncio.run.side_effect = mock_run
+        mock_asyncio_run.assert_called_once()
-
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
-        command = AddCommand()
+        mock_cognee.add.assert_awaited_once_with(
-        args = argparse.Namespace(data=["test1.txt", "test2.txt"], dataset_name="test_dataset")
+            data=["test1.txt", "test2.txt"], dataset_name="test_dataset"
-
+        )
-        command.execute(args)
-
-        mock_asyncio.run.assert_called_once()
 
     @patch("cognee.cli.commands.add_command.asyncio.run")
     def test_execute_with_exception(self, mock_asyncio_run):
@@ -134,9 +126,8 @@ class TestSearchCommand:
         assert actions["top_k"].default == 10
         assert actions["output_format"].default == "pretty"
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.search_command.asyncio")
+    def test_execute_basic_search(self, mock_asyncio_run):
-    def test_execute_basic_search(self, mock_asyncio, mock_import):
         """Test execute with basic search"""
         # Mock the cognee module and SearchType
         mock_cognee = MagicMock()
@@ -144,40 +135,30 @@ class TestSearchCommand:
         mock_search_type = MagicMock()
         mock_search_type.__getitem__.return_value = "GRAPH_COMPLETION"
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-            if name == "cognee":
+            command = SearchCommand()
-                return mock_cognee
+            args = argparse.Namespace(
-            elif name == "cognee.modules.search.types":
+                query_text="test query",
-                module = MagicMock()
+                query_type="GRAPH_COMPLETION",
-                module.SearchType = mock_search_type
+                datasets=None,
-                return module
+                top_k=10,
-            return MagicMock()
+                system_prompt=None,
+                output_format="pretty",
+            )
+            command.execute(args)
 
-        mock_import.side_effect = mock_import_func
+        mock_asyncio_run.assert_called_once()
-
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
-        # Mock asyncio.run to properly handle coroutines
+        mock_cognee.search.assert_awaited_once_with(
-        def mock_run(coro):
-            loop = asyncio.new_event_loop()
-            try:
-                return loop.run_until_complete(coro)
-            finally:
-                loop.close()
-
-        mock_asyncio.run.side_effect = mock_run
-
-        command = SearchCommand()
-        args = argparse.Namespace(
             query_text="test query",
-            query_type="GRAPH_COMPLETION",
+            query_type=ANY,
             datasets=None,
             top_k=10,
-            system_prompt=None,
+            system_prompt_path="answer_simple_question.txt",
-            output_format="pretty",
         )
-
+        # verify the enum’s name separately
-        command.execute(args)
+        called_enum = mock_cognee.search.await_args.kwargs["query_type"]
-
+        assert called_enum.name == "GRAPH_COMPLETION"
-        mock_asyncio.run.assert_called_once()
 
     @patch("cognee.cli.commands.search_command.asyncio.run")
     def test_execute_with_exception(self, mock_asyncio_run):
@@ -227,49 +208,37 @@ class TestCognifyCommand:
         # Check default values
         assert actions["chunker"].default == "TextChunker"
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.cognify_command.asyncio")
+    def test_execute_basic_cognify(self, mock_asyncio_run):
-    def test_execute_basic_cognify(self, mock_asyncio, mock_import):
         """Test execute with basic cognify"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.cognify = AsyncMock(return_value="success")
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-            if name == "cognee":
+            command = CognifyCommand()
-                return mock_cognee
+            args = argparse.Namespace(
-            elif name == "cognee.modules.chunking":
+                datasets=None,
-                module = MagicMock()
+                chunk_size=None,
-                module.TextChunker = MagicMock()
+                ontology_file=None,
-                return module
+                chunker="TextChunker",
-            return MagicMock()
+                background=False,
+                verbose=False,
+            )
+            command.execute(args)
 
-        mock_import.side_effect = mock_import_func
+        mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        from cognee.modules.chunking import TextChunker
 
-        # Mock asyncio.run to properly handle coroutines
+        mock_cognee.cognify.assert_awaited_once_with(
-        def mock_run(coro):
-            loop = asyncio.new_event_loop()
-            try:
-                return loop.run_until_complete(coro)
-            finally:
-                loop.close()
-
-        mock_asyncio.run.side_effect = mock_run
-
-        command = CognifyCommand()
-        args = argparse.Namespace(
             datasets=None,
             chunk_size=None,
-            ontology_file=None,
+            ontology_file_path=None,
-            chunker="TextChunker",
+            chunker=TextChunker,
-            background=False,
+            run_in_background=False,
-            verbose=False,
         )
 
-        command.execute(args)
-
-        mock_asyncio.run.assert_called_once()
-
     @patch("cognee.cli.commands.cognify_command.asyncio.run")
     def test_execute_with_exception(self, mock_asyncio_run):
         """Test execute handles exceptions properly"""
@@ -314,36 +283,27 @@ class TestDeleteCommand:
         assert "force" in actions
 
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.delete_command.asyncio")
+    def test_execute_delete_dataset_with_confirmation(self, mock_asyncio_run, mock_confirm):
-    def test_execute_delete_dataset_with_confirmation(
-        self, mock_asyncio, mock_import, mock_confirm
-    ):
         """Test execute delete dataset with user confirmation"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.delete = AsyncMock()
-        mock_import.return_value = mock_cognee
 
-        # Mock asyncio.run to properly handle coroutines
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        def mock_run(coro):
+            command = DeleteCommand()
-            loop = asyncio.new_event_loop()
+            args = argparse.Namespace(
-            try:
+                dataset_name="test_dataset", user_id=None, all=False, force=False
-                return loop.run_until_complete(coro)
+            )
-            finally:
-                loop.close()
 
-        mock_asyncio.run.side_effect = mock_run
+            mock_confirm.return_value = True
 
-        command = DeleteCommand()
+            command.execute(args)
-        args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=False)
 
-        mock_confirm.return_value = True
+        mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?")
-
+        mock_asyncio_run.assert_called_once()
-        command.execute(args)
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
-
+        mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None)
-        mock_confirm.assert_called_once()
-        mock_asyncio.run.assert_called_once()
 
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
     def test_execute_delete_cancelled(self, mock_confirm):
@@ -356,33 +316,26 @@ class TestDeleteCommand:
         # Should not raise exception, just return
         command.execute(args)
 
-        mock_confirm.assert_called_once()
+        mock_confirm.assert_called_once_with(f"Delete dataset '{args.dataset_name}'?")
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.delete_command.asyncio")
+    def test_execute_delete_forced(self, mock_asyncio_run):
-    def test_execute_delete_forced(self, mock_asyncio, mock_import):
         """Test execute delete with force flag"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.delete = AsyncMock()
-        mock_import.return_value = mock_cognee
 
-        # Mock asyncio.run to properly handle coroutines
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        def mock_run(coro):
+            command = DeleteCommand()
-            loop = asyncio.new_event_loop()
+            args = argparse.Namespace(
-            try:
+                dataset_name="test_dataset", user_id=None, all=False, force=True
-                return loop.run_until_complete(coro)
+            )
-            finally:
-                loop.close()
 
-        mock_asyncio.run.side_effect = mock_run
+            command.execute(args)
 
-        command = DeleteCommand()
+        mock_asyncio_run.assert_called_once()
-        args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=True)
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
-
+        mock_cognee.delete.assert_awaited_once_with(dataset_name="test_dataset", user_id=None)
-        command.execute(args)
-
-        mock_asyncio.run.assert_called_once()
 
     def test_execute_no_delete_target(self):
         """Test execute when no delete target is specified"""
@@ -490,20 +443,19 @@ class TestConfigCommand:
         command.execute(args)
 
     @patch("cognee.cli.commands.config_command.fmt.confirm")
-    @patch("builtins.__import__")
+    def test_execute_unset_action(self, mock_confirm):
-    def test_execute_unset_action(self, mock_import, mock_confirm):
         """Test execute unset action"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.config.set_llm_provider = MagicMock()
-        mock_import.return_value = mock_cognee
 
-        command = ConfigCommand()
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        args = argparse.Namespace(config_action="unset", key="llm_provider", force=False)
+            command = ConfigCommand()
+            args = argparse.Namespace(config_action="unset", key="llm_provider", force=False)
 
-        mock_confirm.return_value = True
+            mock_confirm.return_value = True
 
-        command.execute(args)
+            command.execute(args)
 
         mock_confirm.assert_called_once()
 

cognee/tests/unit/cli/test_cli_edge_cases.py

@@ -3,9 +3,10 @@ Tests for CLI edge cases and error scenarios with proper mocking.
 """
 
 import pytest
+import sys
+import asyncio
 import argparse
-import types
+from unittest.mock import patch, MagicMock, AsyncMock, ANY, call
-from unittest.mock import patch, MagicMock, AsyncMock
 from cognee.cli.commands.add_command import AddCommand
 from cognee.cli.commands.search_command import SearchCommand
 from cognee.cli.commands.cognify_command import CognifyCommand
@@ -14,27 +15,32 @@ 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 TestAddCommandEdgeCases:
     """Test edge cases for AddCommand"""
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.add_command.asyncio.run")
+    def test_add_empty_data_list(self, mock_asyncio_run):
-    def test_add_empty_data_list(self, mock_asyncio_run, mock_import):
-        """Test add command with empty data list"""
-        # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.add = AsyncMock()
-        mock_import.return_value = mock_cognee
 
-        command = AddCommand()
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        # This shouldn't happen due to argparse nargs="+", but test defensive coding
+            command = AddCommand()
-        args = argparse.Namespace(data=[], dataset_name="test_dataset")
+            args = argparse.Namespace(data=[], dataset_name="test_dataset")
+            command.execute(args)
 
-        command.execute(args)
         mock_asyncio_run.assert_called_once()
-
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
-        coro = mock_asyncio_run.call_args[0][0]
+        mock_cognee.add.assert_awaited_once_with(data=[], dataset_name="test_dataset")
-        assert isinstance(coro, types.CoroutineType)
 
     @patch("cognee.cli.commands.add_command.asyncio.run")
     def test_add_asyncio_run_exception(self, mock_asyncio_run):
@@ -71,9 +77,8 @@ class TestAddCommandEdgeCases:
 class TestSearchCommandEdgeCases:
     """Test edge cases for SearchCommand"""
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.search_command.asyncio.run")
+    def test_search_empty_results(self, mock_asyncio_run):
-    def test_search_empty_results(self, mock_asyncio_run, mock_import):
         """Test search command with empty results"""
         # Mock the cognee module and SearchType
         mock_cognee = MagicMock()
@@ -81,35 +86,36 @@ class TestSearchCommandEdgeCases:
         mock_search_type = MagicMock()
         mock_search_type.__getitem__.return_value = "GRAPH_COMPLETION"
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-            if name == "cognee":
+            command = SearchCommand()
-                return mock_cognee
+            args = argparse.Namespace(
-            elif name == "cognee.modules.search.types":
+                query_text="nonexistent query",
-                module = MagicMock()
+                query_type="GRAPH_COMPLETION",
-                module.SearchType = mock_search_type
+                datasets=None,
-                return module
+                top_k=10,
-            return MagicMock()
+                system_prompt=None,
+                output_format="pretty",
+            )
+
+            # Should handle empty results gracefully
+            command.execute(args)
 
-        mock_import.side_effect = mock_import_func
         mock_asyncio_run.return_value = []
-
+        mock_asyncio_run.assert_called_once()
-        command = SearchCommand()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
-        args = argparse.Namespace(
+        mock_cognee.search.assert_awaited_once_with(
             query_text="nonexistent query",
-            query_type="GRAPH_COMPLETION",
+            query_type=ANY,
             datasets=None,
             top_k=10,
-            system_prompt=None,
+            system_prompt_path="answer_simple_question.txt",
-            output_format="pretty",
         )
+        # verify the enum’s name separately
+        called_enum = mock_cognee.search.await_args.kwargs["query_type"]
+        assert called_enum.name == "GRAPH_COMPLETION"
 
-        # Should handle empty results gracefully
+    @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run)
-        command.execute(args)
+    def test_search_very_large_top_k(self, mock_asyncio_run):
-        mock_asyncio_run.assert_called_once()
-
-    @patch("builtins.__import__")
-    @patch("cognee.cli.commands.search_command.asyncio.run")
-    def test_search_very_large_top_k(self, mock_asyncio_run, mock_import):
         """Test search command with very large top-k value"""
         # Mock the cognee module and SearchType
         mock_cognee = MagicMock()
@@ -117,30 +123,33 @@ class TestSearchCommandEdgeCases:
         mock_search_type = MagicMock()
         mock_search_type.__getitem__.return_value = "CHUNKS"
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
-            if name == "cognee":
-                return mock_cognee
-            elif name == "cognee.modules.search.types":
-                module = MagicMock()
-                module.SearchType = mock_search_type
-                return module
-            return MagicMock()
-
-        mock_import.side_effect = mock_import_func
         mock_asyncio_run.return_value = ["result1"]
 
-        command = SearchCommand()
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        args = argparse.Namespace(
+            command = SearchCommand()
-            query_text="test query",
+            args = argparse.Namespace(
-            query_type="CHUNKS",
+                query_text="test query",
-            datasets=None,
+                query_type="CHUNKS",
-            top_k=999999,  # Very large value
+                datasets=None,
-            system_prompt=None,
+                top_k=999999,  # Very large value
-            output_format="json",
+                system_prompt=None,
-        )
+                output_format="json",
+            )
+
+            command.execute(args)
 
-        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=999999,
+            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 == "CHUNKS"
 
     @patch("builtins.__import__")
     def test_search_invalid_search_type_enum(self, mock_import):
@@ -181,9 +190,8 @@ class TestSearchCommandEdgeCases:
         args = parser.parse_args([unicode_query])
         assert args.query_text == unicode_query
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.search_command.asyncio.run")
+    def test_search_results_with_none_values(self, mock_asyncio_run):
-    def test_search_results_with_none_values(self, mock_asyncio_run, mock_import):
         """Test search command when results contain None values"""
         # Mock the cognee module and SearchType
         mock_cognee = MagicMock()
@@ -191,105 +199,105 @@ class TestSearchCommandEdgeCases:
         mock_search_type = MagicMock()
         mock_search_type.__getitem__.return_value = "CHUNKS"
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-            if name == "cognee":
+            command = SearchCommand()
-                return mock_cognee
+            args = argparse.Namespace(
-            elif name == "cognee.modules.search.types":
+                query_text="test query",
-                module = MagicMock()
+                query_type="CHUNKS",
-                module.SearchType = mock_search_type
+                datasets=None,
-                return module
+                top_k=10,
-            return MagicMock()
+                system_prompt=None,
+                output_format="pretty",
+            )
 
-        mock_import.side_effect = mock_import_func
+            # Should handle None values gracefully
-        mock_asyncio_run.return_value = [None, "valid result", None]
+            command.execute(args)
 
-        command = SearchCommand()
+        mock_asyncio_run.assert_called_once()
-        args = argparse.Namespace(
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
+        mock_cognee.search.assert_awaited_once_with(
             query_text="test query",
-            query_type="CHUNKS",
+            query_type=ANY,
             datasets=None,
             top_k=10,
-            system_prompt=None,
+            system_prompt_path="answer_simple_question.txt",
-            output_format="pretty",
         )
-
+        # verify the enum’s name separately
-        # Should handle None values gracefully
+        called_enum = mock_cognee.search.await_args.kwargs["query_type"]
-        command.execute(args)
+        assert called_enum.name == "CHUNKS"
-        mock_asyncio_run.assert_called_once()
 
 
 class TestCognifyCommandEdgeCases:
     """Test edge cases for CognifyCommand"""
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.cognify_command.asyncio.run")
+    def test_cognify_invalid_chunk_size(self, mock_asyncio_run):
-    def test_cognify_invalid_chunk_size(self, mock_asyncio_run, mock_import):
         """Test cognify command with invalid chunk size"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.cognify = AsyncMock()
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-            if name == "cognee":
+            command = CognifyCommand()
-                return mock_cognee
+            args = argparse.Namespace(
-            elif name == "cognee.modules.chunking":
+                datasets=None,
-                module = MagicMock()
+                chunk_size=-100,  # Invalid negative chunk size
-                module.TextChunker = MagicMock()
+                ontology_file=None,
-                return module
+                chunker="TextChunker",
-            return MagicMock()
+                background=False,
+                verbose=False,
+            )
 
-        mock_import.side_effect = mock_import_func
+            # Should pass the invalid value to cognify and let it handle the validation
+            command.execute(args)
 
-        command = CognifyCommand()
+        mock_asyncio_run.assert_called_once()
-        args = argparse.Namespace(
+        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=-100,  # Invalid negative chunk size
+            chunk_size=-100,
-            ontology_file=None,
+            ontology_file_path=None,
-            chunker="TextChunker",
+            chunker=TextChunker,
-            background=False,
+            run_in_background=False,
-            verbose=False,
         )
 
-        # Should pass the invalid value to cognify and let it handle the validation
+    @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run)
-        command.execute(args)
+    def test_cognify_nonexistent_ontology_file(self, mock_asyncio_run):
-        mock_asyncio_run.assert_called_once()
-
-    @patch("builtins.__import__")
-    @patch("cognee.cli.commands.cognify_command.asyncio.run")
-    def test_cognify_nonexistent_ontology_file(self, mock_asyncio_run, mock_import):
         """Test cognify command with nonexistent ontology file"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.cognify = AsyncMock()
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-            if name == "cognee":
+            command = CognifyCommand()
-                return mock_cognee
+            args = argparse.Namespace(
-            elif name == "cognee.modules.chunking":
+                datasets=None,
-                module = MagicMock()
+                chunk_size=None,
-                module.TextChunker = MagicMock()
+                ontology_file="/nonexistent/path/ontology.owl",
-                return module
+                chunker="TextChunker",
-            return MagicMock()
+                background=False,
+                verbose=False,
+            )
 
-        mock_import.side_effect = mock_import_func
+            # Should pass the path to cognify and let it handle file validation
+            command.execute(args)
 
-        command = CognifyCommand()
+        mock_asyncio_run.assert_called_once()
-        args = argparse.Namespace(
+        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="/nonexistent/path/ontology.owl",
+            ontology_file_path="/nonexistent/path/ontology.owl",
-            chunker="TextChunker",
+            chunker=TextChunker,
-            background=False,
+            run_in_background=False,
-            verbose=False,
         )
 
-        # Should pass the path to cognify and let it handle file validation
-        command.execute(args)
-        mock_asyncio_run.assert_called_once()
-
-    @patch("builtins.__import__")
     @patch("cognee.cli.commands.cognify_command.asyncio.run")
-    def test_cognify_langchain_chunker_import_error(self, mock_asyncio_run, mock_import):
+    def test_cognify_langchain_chunker_import_error(self, mock_asyncio_run):
         """Test cognify command when LangchainChunker import fails"""
         # Mock the cognee module
         mock_cognee = MagicMock()
@@ -306,76 +314,83 @@ class TestCognifyCommandEdgeCases:
                 return module
             return MagicMock()
 
-        mock_import.side_effect = mock_import_func
+        with (
+            patch("builtins.__import__", side_effect=mock_import_func),
+            patch.dict(sys.modules, {"cognee": mock_cognee}),
+        ):
+            command = CognifyCommand()
+            args = argparse.Namespace(
+                datasets=None,
+                chunk_size=None,
+                ontology_file=None,
+                chunker="LangchainChunker",
+                background=False,
+                verbose=True,
+            )
 
-        command = CognifyCommand()
+            # Should fall back to TextChunker and show warning
-        args = argparse.Namespace(
+            command.execute(args)
-            datasets=None,
-            chunk_size=None,
-            ontology_file=None,
-            chunker="LangchainChunker",
-            background=False,
-            verbose=True,
-        )
 
-        # Should fall back to TextChunker and show warning
-        command.execute(args)
         mock_asyncio_run.assert_called_once()
+        assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0])
 
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.cognify_command.asyncio.run")
+    def test_cognify_empty_datasets_list(self, mock_asyncio_run):
-    def test_cognify_empty_datasets_list(self, mock_asyncio_run, mock_import):
         """Test cognify command with nonexistent ontology file"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.cognify = AsyncMock()
 
-        def mock_import_func(name, fromlist=None, *args, **kwargs):
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-            if name == "cognee":
+            command = CognifyCommand()
-                return mock_cognee
+            args = argparse.Namespace(
-            elif name == "cognee.modules.chunking":
+                datasets=[],
-                module = MagicMock()
+                chunk_size=None,
-                module.TextChunker = MagicMock()
+                ontology_file=None,
-                return module
+                chunker="TextChunker",
-            return MagicMock()
+                background=False,
+                verbose=False,
+            )
 
-        mock_import.side_effect = mock_import_func
+            command.execute(args)
 
-        command = CognifyCommand()
-        args = argparse.Namespace(
-            datasets=[],
-            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,
+        )
 
 
 class TestDeleteCommandEdgeCases:
     """Test edge cases for DeleteCommand"""
 
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
-    @patch("builtins.__import__")
+    @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
-    @patch("cognee.cli.commands.delete_command.asyncio.run")
+    def test_delete_all_with_user_id(self, mock_asyncio_run, mock_confirm):
-    def test_delete_all_with_user_id(self, mock_asyncio_run, mock_import, mock_confirm):
         """Test delete command with both --all and --user-id"""
         # Mock the cognee module
         mock_cognee = MagicMock()
         mock_cognee.delete = AsyncMock()
-        mock_import.return_value = mock_cognee
 
-        command = DeleteCommand()
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        args = argparse.Namespace(dataset_name=None, user_id="test_user", all=True, force=False)
+            command = DeleteCommand()
+            args = argparse.Namespace(dataset_name=None, user_id="test_user", all=True, force=False)
 
-        mock_confirm.return_value = True
+            mock_confirm.return_value = True
 
-        # Should handle both flags being set
+            # Should handle both flags being set
-        command.execute(args)
+            command.execute(args)
+
+        mock_confirm.assert_called_once_with("Delete ALL data from cognee?")
         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=None, user_id="test_user")
 
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
     def test_delete_confirmation_keyboard_interrupt(self, mock_confirm):
@@ -482,20 +497,20 @@ class TestConfigCommandEdgeCases:
         mock_cognee.config.set.assert_called_once_with("test_key", invalid_json)
 
     @patch("cognee.cli.commands.config_command.fmt.confirm")
-    @patch("builtins.__import__")
+    def test_config_unset_unknown_key(self, mock_confirm):
-    def test_config_unset_unknown_key(self, mock_import, mock_confirm):
         """Test config unset with unknown key"""
         # Mock the cognee module
         mock_cognee = MagicMock()
-        mock_import.return_value = mock_cognee
 
-        command = ConfigCommand()
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
-        args = argparse.Namespace(config_action="unset", key="unknown_key", force=False)
+            command = ConfigCommand()
+            args = argparse.Namespace(config_action="unset", key="unknown_key", force=False)
 
-        mock_confirm.return_value = True
+            mock_confirm.return_value = True
+
+            # Should show error for unknown key
+            command.execute(args)
 
-        # Should show error for unknown key
-        command.execute(args)
         mock_confirm.assert_called_once()
 
     @patch("builtins.__import__")