fix: set context properly before running cognee methods tests

This commit is contained in:
Boris Arzentar 2025-11-14 12:36:32 +01:00
parent 0c2eeb873b
commit 9315f3feb5
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
3 changed files with 13 additions and 3 deletions

View file

@ -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={

View file

@ -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)

View file

@ -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={