DataOps-platform/tests/knowledge/test_lightrag_retriever.py

from __future__ import annotations


def test_lightrag_context_without_active_canonical_mapping_is_discarded():
    from app.core.knowledge.access import KnowledgeAccessContext
    from app.core.knowledge.retrieval.lightrag import LightRAGRetriever

    class Client:
        def query_context(self, *_args, **_kwargs):
            return "[DATAOPS_SOURCE DataFlow:removed-object:1] old context"

    context = KnowledgeAccessContext(
        subject_id="user-1",
        roles=frozenset({"viewer"}),
        permissions=frozenset({"governance:read"}),
        business_domain_uids=frozenset({"domain-a"}),
        correlation_id="correlation-1",
    )
    retriever = LightRAGRetriever(Client(), canonical_loader=lambda *_args: ())

    assert retriever.retrieve("query", context, 10) == ()


def test_lightrag_context_is_always_reauthorized_by_canonical_loader():
    from app.core.knowledge.access import KnowledgeAccessContext
    from app.core.knowledge.retrieval.lightrag import LightRAGRetriever

    class Client:
        def query_context(self, *_args, **_kwargs):
            return "[DATAOPS_SOURCE DataFlow:object-1:2] active context"

    captured = []

    def loader(external_ids, context, limit):
        captured.append((external_ids, context.subject_id, limit))
        return ()

    context = KnowledgeAccessContext(
        subject_id="user-1",
        roles=frozenset({"viewer"}),
        permissions=frozenset({"governance:read"}),
        business_domain_uids=frozenset({"domain-a"}),
        correlation_id="correlation-1",
    )
    LightRAGRetriever(Client(), canonical_loader=loader).retrieve("query", context, 10)

    assert captured == [(("DataFlow:object-1:2",), "user-1", 10)]