fix: get_graph_from_model added nodes and edges default values
This commit is contained in:
parent
3f0eb3b41f
commit
fdc6113d11
4 changed files with 47 additions and 3 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue