fix: get_graph_from_model added nodes and edges default values

This commit is contained in:
Boris Arzentar 2025-10-14 21:55:30 +02:00
parent 3f0eb3b41f
commit fdc6113d11
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
4 changed files with 47 additions and 3 deletions

View file

@ -152,10 +152,10 @@ async def get_graph_from_model(
Returns:
Tuple of (nodes, edges) extracted from the model
"""
if not added_nodes:
if added_nodes == None:
added_nodes = {}
if not added_edges:
if added_edges == None:
added_edges = {}
if str(data_point.id) in added_nodes:

View file

@ -2,12 +2,14 @@
Tests for individual CLI commands with proper mocking and coroutine handling.
"""
import os
import pytest
import sys
import argparse
import asyncio
from uuid import uuid4
from unittest.mock import patch, MagicMock, AsyncMock, ANY
import cognee
from cognee.cli.commands.add_command import AddCommand
from cognee.cli.commands.search_command import SearchCommand
from cognee.cli.commands.cognify_command import CognifyCommand
@ -15,6 +17,7 @@ from cognee.cli.commands.delete_command import DeleteCommand
from cognee.cli.commands.config_command import ConfigCommand
from cognee.cli.exceptions import CliCommandException
from cognee.modules.data.methods.get_deletion_counts import DeletionCountsPreview
from cognee.modules.engine.operations.setup import setup
from cognee.modules.users.models import User
@ -289,10 +292,26 @@ class TestDeleteCommand:
@patch("cognee.cli.commands.delete_command.cognee_datasets.delete_dataset")
@patch("cognee.cli.commands.delete_command.fmt.confirm")
@patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
# @pytest.mark.asyncio
def test_execute_delete_dataset_with_confirmation(
self, mock_asyncio_run, mock_confirm, delete_dataset_mock, get_user_mock
):
"""Test execute delete dataset with user confirmation"""
data_directory_path = os.path.join(
os.path.dirname(__file__), ".data_storage/test_cli_commands"
)
cognee_directory_path = os.path.join(
os.path.dirname(__file__), ".cognee_system/test_cli_commands"
)
cognee.config.data_root_directory(data_directory_path)
cognee.config.system_root_directory(cognee_directory_path)
asyncio.run(cognee.prune.prune_data())
asyncio.run(cognee.prune.prune_system(metadata=True))
asyncio.run(setup())
expected_user_id = uuid4()
expected_dataset_id = uuid4()
@ -309,6 +328,9 @@ class TestDeleteCommand:
dataset_id=expected_dataset_id, user_id=expected_user_id
)
asyncio.run(cognee.prune.prune_data())
asyncio.run(cognee.prune.prune_system(metadata=True))
@patch("cognee.cli.commands.delete_command.get_deletion_counts")
@patch("cognee.cli.commands.delete_command.fmt.confirm")
def test_execute_delete_cancelled(self, mock_confirm, mock_get_deletion_counts):

View file

@ -2,12 +2,15 @@
Tests for CLI edge cases and error scenarios with proper mocking.
"""
import os
import pytest
import sys
import asyncio
import argparse
from uuid import uuid4
from unittest.mock import patch, MagicMock, AsyncMock, ANY
import cognee
from cognee.cli.commands.add_command import AddCommand
from cognee.cli.commands.search_command import SearchCommand
from cognee.cli.commands.cognify_command import CognifyCommand
@ -15,6 +18,7 @@ from cognee.cli.commands.delete_command import DeleteCommand
from cognee.cli.commands.config_command import ConfigCommand
from cognee.cli.exceptions import CliCommandException
from cognee.modules.data.methods.get_deletion_counts import DeletionCountsPreview
from cognee.modules.engine.operations.setup import setup
# Mock asyncio.run to properly handle coroutines
@ -385,6 +389,21 @@ class TestDeleteCommandEdgeCases:
@patch("cognee.cli.commands.delete_command.fmt.confirm")
def test_delete_all_with_user_id(self, fmt_confirm_mock, delete_all_mock, async_run_mock):
"""Test delete command with both --all and --user-id"""
data_directory_path = os.path.join(
os.path.dirname(__file__), ".data_storage/test_cli_commands"
)
cognee_directory_path = os.path.join(
os.path.dirname(__file__), ".cognee_system/test_cli_commands"
)
cognee.config.data_root_directory(data_directory_path)
cognee.config.system_root_directory(cognee_directory_path)
asyncio.run(cognee.prune.prune_data())
asyncio.run(cognee.prune.prune_system(metadata=True))
asyncio.run(setup())
fmt_confirm_mock.return_value = True
expected_user_id = uuid4()
@ -399,6 +418,9 @@ class TestDeleteCommandEdgeCases:
delete_all_mock.assert_called_once_with(user_id=expected_user_id)
asyncio.run(cognee.prune.prune_data())
asyncio.run(cognee.prune.prune_system(metadata=True))
@patch("cognee.cli.commands.delete_command.get_deletion_counts")
@patch("cognee.cli.commands.delete_command.fmt.confirm")
def test_delete_confirmation_keyboard_interrupt(self, mock_confirm, mock_get_deletion_counts):

View file

@ -90,7 +90,7 @@ async def test_circular_reference_extraction():
code_file.contains.extend(
[
CodePart(
part_of=code_file,
# part_of=code_file,
source_code=f"Part {part_index}",
)
for part_index in range(2)