diff --git a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py index 11209375..bab0f3bb 100644 --- a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py +++ b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py @@ -118,10 +118,10 @@ def add_attribute(self, key: str, value: Any) -> None: def get_attribute(self, key: str) -> Optional[Union[str, int, float]]: return self.attributes.get(key) - def get_node_from(self): + def get_source_node(self): return self.node1 - def get_node_to(self): + def get_destination_node(self): return self.node2 def __repr__(self) -> str: diff --git a/cognee/modules/retrieval/description_to_codepart_search.py b/cognee/modules/retrieval/description_to_codepart_search.py index 6d5b46b4..e1da9a43 100644 --- a/cognee/modules/retrieval/description_to_codepart_search.py +++ b/cognee/modules/retrieval/description_to_codepart_search.py @@ -85,7 +85,7 @@ async def code_description_to_code_part( for code_file in node_to_search_from.get_skeleton_neighbours(): for code_file_edge in code_file.get_skeleton_edges(): if code_file_edge.get_attribute('relationship_name') == 'contains': - code_pieces_to_return.add(code_file_edge.get_node_to()) + code_pieces_to_return.add(code_file_edge.get_destination_node()) logging.info("Search completed for user: %s, query: '%s'. Found %d code pieces.", user.id, query, len(code_pieces_to_return)) diff --git a/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py b/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py index cc955b0c..4c39883a 100644 --- a/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py +++ b/cognee/tests/unit/modules/retriever/test_description_to_codepart_search.py @@ -6,6 +6,7 @@ @pytest.mark.asyncio async def test_code_description_to_code_part_no_results(): """Test that code_description_to_code_part handles no search results.""" + mock_user = AsyncMock() mock_user.id = "user123" mock_vector_engine = AsyncMock() @@ -24,7 +25,9 @@ async def test_code_description_to_code_part_no_results(): @pytest.mark.asyncio async def test_code_description_to_code_part_invalid_query(): """Test that code_description_to_code_part raises ValueError for invalid query.""" + mock_user = AsyncMock() + with pytest.raises(ValueError, match="The query must be a non-empty string."): from cognee.modules.retrieval.description_to_codepart_search import code_description_to_code_part await code_description_to_code_part("", mock_user, 2) @@ -33,7 +36,9 @@ async def test_code_description_to_code_part_invalid_query(): @pytest.mark.asyncio async def test_code_description_to_code_part_invalid_top_k(): """Test that code_description_to_code_part raises ValueError for invalid top_k.""" + mock_user = AsyncMock() + with pytest.raises(ValueError, match="top_k must be a positive integer."): from cognee.modules.retrieval.description_to_codepart_search import code_description_to_code_part await code_description_to_code_part("search query", mock_user, 0) @@ -42,7 +47,9 @@ async def test_code_description_to_code_part_invalid_top_k(): @pytest.mark.asyncio async def test_code_description_to_code_part_initialization_error(): """Test that code_description_to_code_part raises RuntimeError for engine initialization errors.""" + mock_user = AsyncMock() + with patch("cognee.modules.retrieval.description_to_codepart_search.get_vector_engine", side_effect=Exception("Engine init failed")), \ patch("cognee.modules.retrieval.description_to_codepart_search.get_graph_engine", return_value=AsyncMock()): @@ -54,6 +61,7 @@ async def test_code_description_to_code_part_initialization_error(): @pytest.mark.asyncio async def test_code_description_to_code_part_execution_error(): """Test that code_description_to_code_part raises RuntimeError for execution errors.""" + mock_user = AsyncMock() mock_user.id = "user123" mock_vector_engine = AsyncMock()