Function test_local_file_deletion

cognee/tests/test_pgvector.py:14–50

(data_text, file_location, dataset_1_id, dataset_2_id)

async def test_local_file_deletion(data_text, file_location, dataset_1_id, dataset_2_id):
    from sqlalchemy import select
    import hashlib
    from cognee.infrastructure.databases.relational import get_relational_engine

    engine = get_relational_engine()

    async with engine.get_async_session() as session:
        # Get hash of data contents
        encoded_text = data_text.encode("utf-8")
        data_hash = hashlib.md5(encoded_text).hexdigest()
        # Get data entry from database based on hash contents
        data = (await session.scalars(select(Data).where(Data.content_hash == data_hash))).one()
        assert os.path.isfile(data.raw_data_location.replace("file://", "")), (
            f"Data location doesn't exist: {data.raw_data_location}"
        )
        # Test deletion of data along with local files created by cognee
        await engine.delete_data_entity(data.id, dataset_id=dataset_2_id)
        assert not os.path.exists(data.raw_data_location.replace("file://", "")), (
            f"Data location still exists after deletion: {data.raw_data_location}"
        )

    async with engine.get_async_session() as session:
        # Get data entry from database based on file path
        data = (
            await session.scalars(
                select(Data).where(Data.original_data_location == "file://" + file_location)
            )
        ).one()
        assert os.path.isfile(data.original_data_location.replace("file://", "")), (
            f"Data location doesn't exist: {data.original_data_location}"
        )
        # Test local files not created by cognee won't get deleted
        await engine.delete_data_entity(data.id, dataset_id=dataset_1_id)
        assert os.path.exists(data.original_data_location.replace("file://", "")), (
            f"Data location doesn't exists: {data.original_data_location}"
        )


async def test_getting_of_documents(dataset_id_1):