gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: unit tests

Browse source
This commit is contained in:
Boris Arzentar 2025-10-30 22:23:01 +01:00
parent ae171877d8
commit 1864fc8f2b
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
4 changed files with 25 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
cognee/modules/pipelines/operations/run_tasks.py
View file

@ -7,7 +7,7 @@ from uuid import UUID
from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.databases.relational import get_relational_engine from cognee.infrastructure.databases.relational import get_relational_engine
from cognee.modules.data.methods import get_dataset from cognee.modules.data.methods import get_authorized_dataset
from cognee.modules.pipelines.operations.run_tasks_distributed import run_tasks_distributed from cognee.modules.pipelines.operations.run_tasks_distributed import run_tasks_distributed
from cognee.modules.users.models import User from cognee.modules.users.models import User
from cognee.shared.logging_utils import get_logger from cognee.shared.logging_utils import get_logger
@ -65,7 +65,7 @@ async def run_tasks(
if not user: if not user:
user = await get_default_user() user = await get_default_user()
dataset = await get_dataset(user.id, dataset_id) dataset = await get_authorized_dataset(user, dataset_id, "write")
pipeline_id = generate_pipeline_id(user.id, dataset.id, pipeline_name) pipeline_id = generate_pipeline_id(user.id, dataset.id, pipeline_name)
pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data) pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data)
pipeline_run_id = pipeline_run.pipeline_run_id pipeline_run_id = pipeline_run.pipeline_run_id

5
cognee/modules/pipelines/operations/run_tasks_distributed.py
View file

@ -6,10 +6,9 @@ except ModuleNotFoundError:
from typing import Any, List, Optional from typing import Any, List, Optional
from uuid import UUID from uuid import UUID
from cognee.modules.data.methods import get_dataset from cognee.modules.data.methods import get_authorized_dataset
from cognee.modules.data.models import Dataset from cognee.modules.data.models import Dataset
from cognee.modules.pipelines.tasks.task import Task from cognee.modules.pipelines.tasks.task import Task
from cognee.infrastructure.databases.relational import get_relational_engine
from cognee.modules.pipelines.models import ( from cognee.modules.pipelines.models import (
PipelineRunStarted, PipelineRunStarted,
PipelineRunCompleted, PipelineRunCompleted,
@ -94,7 +93,7 @@ async def run_tasks_distributed(
if not user: if not user:
user = await get_default_user() user = await get_default_user()
dataset = await get_dataset(user.id, dataset_id) dataset = await get_authorized_dataset(user, dataset_id, "write")
pipeline_id: UUID = generate_pipeline_id(user.id, dataset.id, pipeline_name) pipeline_id: UUID = generate_pipeline_id(user.id, dataset.id, pipeline_name)
pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data) pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset.id, data)
pipeline_run_id: UUID = pipeline_run.pipeline_run_id pipeline_run_id: UUID = pipeline_run.pipeline_run_id

19
cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py
View file

@ -288,13 +288,15 @@ class TestDeleteCommand:
assert "all" in actions assert "all" in actions
assert "force" in actions assert "force" in actions
@patch("cognee.modules.data.methods.get_deletion_counts.get_user") # @patch("cognee.modules.users.methods.get_user")
@patch("cognee.cli.commands.delete_command.cognee_datasets.delete_dataset") @patch("cognee.cli.commands.delete_command.get_user")
@patch("cognee.cli.commands.delete_command.cognee_datasets")
@patch("cognee.cli.commands.delete_command.fmt.confirm") @patch("cognee.cli.commands.delete_command.fmt.confirm")
@patch("cognee.cli.commands.delete_command.get_deletion_counts")
@patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run) @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
# @pytest.mark.asyncio # @pytest.mark.asyncio
def test_execute_delete_dataset_with_confirmation( def test_execute_delete_dataset_with_confirmation(
self, mock_asyncio_run, mock_confirm, delete_dataset_mock, get_user_mock self, mock_asyncio_run, get_deletion_counts_mock, mock_confirm, datasets_mock, get_user_mock
): ):
"""Test execute delete dataset with user confirmation""" """Test execute delete dataset with user confirmation"""
data_directory_path = os.path.join( data_directory_path = os.path.join(
@ -314,8 +316,15 @@ class TestDeleteCommand:
expected_user_id = uuid4() expected_user_id = uuid4()
expected_dataset_id = uuid4() expected_dataset_id = uuid4()
expected_user = User(id=expected_user_id)
get_user_mock.return_value = User(id=expected_user_id) get_user_mock.return_value = expected_user
get_deletion_counts_mock = AsyncMock()
get_deletion_counts_mock.return_value = DeletionCountsPreview()
delete_dataset_mock = AsyncMock()
datasets_mock.delete_dataset = delete_dataset_mock
command = DeleteCommand() command = DeleteCommand()
args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id) args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
@ -325,7 +334,7 @@ class TestDeleteCommand:
command.execute(args) command.execute(args)
delete_dataset_mock.assert_awaited_once_with( delete_dataset_mock.assert_awaited_once_with(
dataset_id=expected_dataset_id, user_id=expected_user_id dataset_id=expected_dataset_id, user=expected_user
) )
asyncio.run(cognee.prune.prune_data()) asyncio.run(cognee.prune.prune_data())

8
cognee/tests/test_feedback_enrichment.py
View file

@ -11,8 +11,11 @@ Tests the complete feedback enrichment pipeline:
import os import os
import pathlib import pathlib
from uuid import UUID, uuid4
from collections import Counter from collections import Counter
from pydantic import BaseModel
import cognee import cognee
from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.modules.pipelines.tasks.task import Task from cognee.modules.pipelines.tasks.task import Task
@ -129,10 +132,13 @@ async def main():
Task(link_enrichments_to_feedback), Task(link_enrichments_to_feedback),
] ]
class EnrichmentData(BaseModel):
id: UUID
await cognee.memify( await cognee.memify(
extraction_tasks=extraction_tasks, extraction_tasks=extraction_tasks,
enrichment_tasks=enrichment_tasks, enrichment_tasks=enrichment_tasks,
data=[{}], data=[EnrichmentData(id=uuid4())],
dataset="feedback_enrichment_test_memify", dataset="feedback_enrichment_test_memify",
) )