From 9315f3feb586acf5043aef8075d3ffe30042d922 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Fri, 14 Nov 2025 12:36:32 +0100 Subject: [PATCH] fix: set context properly before running cognee methods tests --- cognee/tests/test_delete_custom_graph.py | 3 +++ .../test_delete_default_graph_with_legacy_data_1.py | 10 +++++++--- cognee/tests/test_delete_permission.py | 3 +++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/cognee/tests/test_delete_custom_graph.py b/cognee/tests/test_delete_custom_graph.py index 42d7d0d63..a9faed078 100644 --- a/cognee/tests/test_delete_custom_graph.py +++ b/cognee/tests/test_delete_custom_graph.py @@ -6,6 +6,7 @@ from pydantic import BaseModel import cognee from cognee.api.v1.datasets import datasets +from cognee.context_global_variables import set_database_global_context_variables from cognee.infrastructure.engine import DataPoint from cognee.modules.data.methods import create_authorized_dataset from cognee.modules.engine.operations.setup import setup @@ -67,6 +68,8 @@ async def main(): data1 = CustomData(id=uuid4()) data2 = CustomData(id=uuid4()) + await set_database_global_context_variables(dataset.id, dataset.owner_id) + await add_data_points( [person1], context={ diff --git a/cognee/tests/test_delete_default_graph_with_legacy_data_1.py b/cognee/tests/test_delete_default_graph_with_legacy_data_1.py index 79f8deeac..0db824830 100644 --- a/cognee/tests/test_delete_default_graph_with_legacy_data_1.py +++ b/cognee/tests/test_delete_default_graph_with_legacy_data_1.py @@ -1,4 +1,5 @@ import os +import json import pathlib from uuid import NAMESPACE_OID, uuid5 import pytest @@ -247,7 +248,8 @@ async def main(mock_create_structured_output: AsyncMock): for node in initial_nodes: node_data = node[1] - collection_name = node_data["type"] + "_" + node_data["metadata"]["index_fields"][0] + node_metadata = json.loads(node_data["metadata"]) + collection_name = node_data["type"] + "_" + node_metadata["index_fields"][0] if collection_name not in initial_nodes_by_vector_collection: initial_nodes_by_vector_collection[collection_name] = [] initial_nodes_by_vector_collection[collection_name].append(node) @@ -269,7 +271,8 @@ async def main(mock_create_structured_output: AsyncMock): after_delete_nodes_by_vector_collection = {} for node in initial_nodes: node_data = node[1] - collection_name = node_data["type"] + "_" + node_data["metadata"]["index_fields"][0] + node_metadata = json.loads(node_data["metadata"]) + collection_name = node_data["type"] + "_" + node_metadata["index_fields"][0] if collection_name not in after_delete_nodes_by_vector_collection: after_delete_nodes_by_vector_collection[collection_name] = [] after_delete_nodes_by_vector_collection[collection_name].append(node) @@ -292,7 +295,8 @@ async def main(mock_create_structured_output: AsyncMock): old_nodes_by_vector_collection = {} for node in old_nodes: - collection_name = node.type + "_" + node.metadata["index_fields"][0] + node_metadata = node.metadata + collection_name = node.type + "_" + node_metadata["index_fields"][0] if collection_name not in old_nodes_by_vector_collection: old_nodes_by_vector_collection[collection_name] = [] old_nodes_by_vector_collection[collection_name].append(node) diff --git a/cognee/tests/test_delete_permission.py b/cognee/tests/test_delete_permission.py index 0735b6f75..dcb2d3d3f 100644 --- a/cognee/tests/test_delete_permission.py +++ b/cognee/tests/test_delete_permission.py @@ -6,6 +6,7 @@ from pydantic import BaseModel import cognee from cognee.api.v1.datasets import datasets +from cognee.context_global_variables import set_database_global_context_variables from cognee.infrastructure.engine import DataPoint from cognee.modules.data.exceptions.exceptions import UnauthorizedDataAccessError from cognee.modules.data.methods import create_authorized_dataset @@ -70,6 +71,8 @@ async def main(): data1 = CustomData(id=uuid4()) data2 = CustomData(id=uuid4()) + await set_database_global_context_variables(dataset.id, dataset.owner_id) + await add_data_points( [person1], context={