From c013499f19ad58229f738c20b95acd5dc011e4b9 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Tue, 26 Aug 2025 18:34:38 +0200 Subject: [PATCH] Fix tests failing for python 3.10 --- cognee/tests/unit/cli/test_cli_commands.py | 256 +++++-------- cognee/tests/unit/cli/test_cli_edge_cases.py | 377 ++++++++++--------- 2 files changed, 300 insertions(+), 333 deletions(-) diff --git a/cognee/tests/unit/cli/test_cli_commands.py b/cognee/tests/unit/cli/test_cli_commands.py index de45b83e4..ed0ca14d1 100644 --- a/cognee/tests/unit/cli/test_cli_commands.py +++ b/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") - def test_execute_single_item(self, mock_asyncio, mock_import): + @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() - mock_import.return_value = mock_cognee - # 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() + 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.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() - args = argparse.Namespace(data=["test.txt"], dataset_name="test_dataset") - - 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): + @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() - mock_import.return_value = mock_cognee - # Mock asyncio.run to properly handle coroutines - def mock_run(coro): - loop = asyncio.new_event_loop() - try: - return loop.run_until_complete(coro) - finally: - loop.close() + 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.side_effect = mock_run - - command = AddCommand() - args = argparse.Namespace(data=["test1.txt", "test2.txt"], dataset_name="test_dataset") - - command.execute(args) - - mock_asyncio.run.assert_called_once() + 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): @@ -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") - def test_execute_basic_search(self, mock_asyncio, mock_import): + @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() @@ -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): - if name == "cognee": - return mock_cognee - elif name == "cognee.modules.search.types": - module = MagicMock() - module.SearchType = mock_search_type - return module - return MagicMock() + 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_import.side_effect = mock_import_func - - # Mock asyncio.run to properly handle coroutines - 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( + 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="GRAPH_COMPLETION", + query_type=ANY, datasets=None, top_k=10, - system_prompt=None, - output_format="pretty", + system_prompt_path="answer_simple_question.txt", ) - - command.execute(args) - - mock_asyncio.run.assert_called_once() + # 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): @@ -227,49 +208,37 @@ class TestCognifyCommand: # Check default values assert actions["chunker"].default == "TextChunker" - @patch("builtins.__import__") - @patch("cognee.cli.commands.cognify_command.asyncio") - def test_execute_basic_cognify(self, mock_asyncio, mock_import): + @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") - def mock_import_func(name, fromlist=None, *args, **kwargs): - if name == "cognee": - return mock_cognee - elif name == "cognee.modules.chunking": - module = MagicMock() - module.TextChunker = MagicMock() - return module - return MagicMock() + 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_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 - 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( + mock_cognee.cognify.assert_awaited_once_with( datasets=None, chunk_size=None, - ontology_file=None, - chunker="TextChunker", - background=False, - verbose=False, + ontology_file_path=None, + chunker=TextChunker, + run_in_background=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") - def test_execute_delete_dataset_with_confirmation( - self, mock_asyncio, mock_import, mock_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() - mock_import.return_value = mock_cognee - # Mock asyncio.run to properly handle coroutines - def mock_run(coro): - loop = asyncio.new_event_loop() - try: - return loop.run_until_complete(coro) - finally: - loop.close() + 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_asyncio.run.side_effect = mock_run + mock_confirm.return_value = True - command = DeleteCommand() - args = argparse.Namespace(dataset_name="test_dataset", user_id=None, all=False, force=False) + command.execute(args) - mock_confirm.return_value = True - - command.execute(args) - - mock_confirm.assert_called_once() - mock_asyncio.run.assert_called_once() + 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): @@ -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") - def test_execute_delete_forced(self, mock_asyncio, mock_import): + @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() - mock_import.return_value = mock_cognee - # Mock asyncio.run to properly handle coroutines - def mock_run(coro): - loop = asyncio.new_event_loop() - try: - return loop.run_until_complete(coro) - finally: - loop.close() + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = DeleteCommand() + args = argparse.Namespace( + dataset_name="test_dataset", user_id=None, all=False, force=True + ) - mock_asyncio.run.side_effect = mock_run + command.execute(args) - 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() + 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""" @@ -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_import, mock_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() - mock_import.return_value = mock_cognee - command = ConfigCommand() - args = argparse.Namespace(config_action="unset", key="llm_provider", force=False) + 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 + mock_confirm.return_value = True - command.execute(args) + command.execute(args) mock_confirm.assert_called_once() diff --git a/cognee/tests/unit/cli/test_cli_edge_cases.py b/cognee/tests/unit/cli/test_cli_edge_cases.py index a35afa81f..0c59accda 100644 --- a/cognee/tests/unit/cli/test_cli_edge_cases.py +++ b/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 +from unittest.mock import patch, MagicMock, AsyncMock, ANY, call 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") - def test_add_empty_data_list(self, mock_asyncio_run, mock_import): - """Test add command with empty data list""" - # Mock the cognee module + @patch("cognee.cli.commands.add_command.asyncio.run", side_effect=_mock_run) + def test_add_empty_data_list(self, mock_asyncio_run): mock_cognee = MagicMock() mock_cognee.add = AsyncMock() - mock_import.return_value = mock_cognee - command = AddCommand() - # This shouldn't happen due to argparse nargs="+", but test defensive coding - args = argparse.Namespace(data=[], dataset_name="test_dataset") + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = AddCommand() + args = argparse.Namespace(data=[], dataset_name="test_dataset") + command.execute(args) - command.execute(args) mock_asyncio_run.assert_called_once() - - coro = mock_asyncio_run.call_args[0][0] - assert isinstance(coro, types.CoroutineType) + assert asyncio.iscoroutine(mock_asyncio_run.call_args[0][0]) + mock_cognee.add.assert_awaited_once_with(data=[], dataset_name="test_dataset") @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") - def test_search_empty_results(self, mock_asyncio_run, mock_import): + @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run) + def test_search_empty_results(self, mock_asyncio_run): """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): - if name == "cognee": - return mock_cognee - elif name == "cognee.modules.search.types": - module = MagicMock() - module.SearchType = mock_search_type - return module - return MagicMock() + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = SearchCommand() + args = argparse.Namespace( + query_text="nonexistent query", + query_type="GRAPH_COMPLETION", + datasets=None, + top_k=10, + 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 = [] - - command = SearchCommand() - args = argparse.Namespace( + 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="nonexistent query", - query_type="GRAPH_COMPLETION", + query_type=ANY, datasets=None, top_k=10, - system_prompt=None, - output_format="pretty", + 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" - # Should handle empty results gracefully - command.execute(args) - 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): + @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run) + def test_search_very_large_top_k(self, mock_asyncio_run): """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() - args = argparse.Namespace( - query_text="test query", - query_type="CHUNKS", - datasets=None, - top_k=999999, # Very large value - system_prompt=None, - output_format="json", - ) + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = SearchCommand() + args = argparse.Namespace( + query_text="test query", + query_type="CHUNKS", + datasets=None, + top_k=999999, # Very large value + 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") - def test_search_results_with_none_values(self, mock_asyncio_run, mock_import): + @patch("cognee.cli.commands.search_command.asyncio.run", side_effect=_mock_run) + def test_search_results_with_none_values(self, mock_asyncio_run): """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): - if name == "cognee": - return mock_cognee - elif name == "cognee.modules.search.types": - module = MagicMock() - module.SearchType = mock_search_type - return module - return MagicMock() + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = SearchCommand() + args = argparse.Namespace( + query_text="test query", + query_type="CHUNKS", + datasets=None, + top_k=10, + system_prompt=None, + output_format="pretty", + ) - mock_import.side_effect = mock_import_func - mock_asyncio_run.return_value = [None, "valid result", None] + # Should handle None values gracefully + command.execute(args) - command = SearchCommand() - args = argparse.Namespace( + 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="CHUNKS", + query_type=ANY, datasets=None, top_k=10, - system_prompt=None, - output_format="pretty", + system_prompt_path="answer_simple_question.txt", ) - - # Should handle None values gracefully - command.execute(args) - mock_asyncio_run.assert_called_once() + # verify the enum’s name separately + called_enum = mock_cognee.search.await_args.kwargs["query_type"] + assert called_enum.name == "CHUNKS" class TestCognifyCommandEdgeCases: """Test edge cases for CognifyCommand""" - @patch("builtins.__import__") - @patch("cognee.cli.commands.cognify_command.asyncio.run") - def test_cognify_invalid_chunk_size(self, mock_asyncio_run, mock_import): + @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run) + def test_cognify_invalid_chunk_size(self, mock_asyncio_run): """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): - if name == "cognee": - return mock_cognee - elif name == "cognee.modules.chunking": - module = MagicMock() - module.TextChunker = MagicMock() - return module - return MagicMock() + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = CognifyCommand() + args = argparse.Namespace( + datasets=None, + chunk_size=-100, # Invalid negative chunk size + ontology_file=None, + chunker="TextChunker", + 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() - args = argparse.Namespace( + 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=-100, # Invalid negative chunk size - ontology_file=None, - chunker="TextChunker", - background=False, - verbose=False, + chunk_size=-100, + ontology_file_path=None, + chunker=TextChunker, + run_in_background=False, ) - # Should pass the invalid value to cognify and let it handle the validation - command.execute(args) - 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): + @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run) + def test_cognify_nonexistent_ontology_file(self, mock_asyncio_run): """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): - if name == "cognee": - return mock_cognee - elif name == "cognee.modules.chunking": - module = MagicMock() - module.TextChunker = MagicMock() - return module - return MagicMock() + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = CognifyCommand() + args = argparse.Namespace( + datasets=None, + chunk_size=None, + ontology_file="/nonexistent/path/ontology.owl", + chunker="TextChunker", + 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() - args = argparse.Namespace( + 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="/nonexistent/path/ontology.owl", - chunker="TextChunker", - background=False, - verbose=False, + ontology_file_path="/nonexistent/path/ontology.owl", + chunker=TextChunker, + run_in_background=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() - args = argparse.Namespace( - 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) - # 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") - def test_cognify_empty_datasets_list(self, mock_asyncio_run, mock_import): + @patch("cognee.cli.commands.cognify_command.asyncio.run", side_effect=_mock_run) + def test_cognify_empty_datasets_list(self, mock_asyncio_run): """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): - if name == "cognee": - return mock_cognee - elif name == "cognee.modules.chunking": - module = MagicMock() - module.TextChunker = MagicMock() - return module - return MagicMock() + with patch.dict(sys.modules, {"cognee": mock_cognee}): + command = CognifyCommand() + args = argparse.Namespace( + datasets=[], + chunk_size=None, + ontology_file=None, + chunker="TextChunker", + 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") - def test_delete_all_with_user_id(self, mock_asyncio_run, mock_import, mock_confirm): + @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) + def test_delete_all_with_user_id(self, mock_asyncio_run, 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() - args = argparse.Namespace(dataset_name=None, user_id="test_user", all=True, force=False) + with patch.dict(sys.modules, {"cognee": mock_cognee}): + 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 - command.execute(args) + # Should handle both flags being set + 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_import, mock_confirm): + def test_config_unset_unknown_key(self, mock_confirm): """Test config unset with unknown key""" # Mock the cognee module mock_cognee = MagicMock() - mock_import.return_value = mock_cognee - command = ConfigCommand() - args = argparse.Namespace(config_action="unset", key="unknown_key", force=False) + with patch.dict(sys.modules, {"cognee": mock_cognee}): + 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__")