DataOps-platform/tests/knowledge/test_lightrag_client.py

from __future__ import annotations

import pytest


class Response:
    def __init__(self, payload, status_code=200):
        self.payload = payload
        self.status_code = status_code

    def raise_for_status(self):
        if self.status_code >= 400:
            raise RuntimeError(f"HTTP {self.status_code}")

    def json(self):
        return self.payload


class Transport:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    def request(self, **kwargs):
        self.calls.append(kwargs)
        response = self.responses.pop(0)
        if isinstance(response, Exception):
            raise response
        return response


def test_lightrag_query_is_context_only_and_workspace_is_not_user_controlled():
    from app.core.knowledge.lightrag.client import LightRAGClient

    transport = Transport([Response({"response": "context text"})])
    client = LightRAGClient(
        base_url="http://lightrag:9621",
        api_key="internal-key",
        workspace="dataops-global-domain-a-g1",
        transport=transport,
    )

    result = client.query_context("上游是什么", mode="mix", limit=10)

    assert result == "context text"
    assert transport.calls[0]["json"]["only_need_context"] is True
    assert transport.calls[0]["json"]["mode"] == "mix"
    assert "workspace" not in transport.calls[0]["json"]


def test_lightrag_circuit_breaker_opens_after_bounded_failures():
    from app.core.knowledge.lightrag.client import (
        LightRAGCircuitOpen,
        LightRAGClient,
    )

    transport = Transport([TimeoutError("one"), TimeoutError("two")])
    client = LightRAGClient(
        base_url="http://lightrag:9621",
        api_key="internal-key",
        workspace="dataops-global-domain-a-g1",
        transport=transport,
        failure_threshold=2,
        cooldown_seconds=60,
    )

    with pytest.raises(TimeoutError):
        client.health()
    with pytest.raises(TimeoutError):
        client.health()
    with pytest.raises(LightRAGCircuitOpen):
        client.health()
    assert len(transport.calls) == 2


def test_lightrag_insert_uses_stable_external_id_as_idempotency_key():
    from app.core.knowledge.lightrag.client import LightRAGClient

    transport = Transport([Response({"track_id": "track-1"})])
    client = LightRAGClient(
        base_url="http://lightrag:9621",
        api_key="internal-key",
        workspace="dataops-global-domain-a-g1",
        transport=transport,
    )

    receipt = client.insert(
        external_document_id="DataFlow:source-1:2",
        content="canonical content",
        metadata={"point_keys": ["DataFlow/source-1/purpose"]},
    )

    assert receipt.track_id == "track-1"
    assert transport.calls[0]["headers"]["Idempotency-Key"] == "DataFlow:source-1:2"