ragflow/test/unit_test/rag/graphrag/test_checkpoint_resume.py

#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

"""Tests for GraphRAG/RAPTOR checkpoint/resume logic.

Calls the real implementations:
  - load_subgraph_from_store  (rag/graphrag/general/index.py)
  - has_raptor_chunks         (rag/svr/task_executor.py)

Both modules are loaded via importlib with their infrastructure dependencies
mocked, so the actual query logic, pagination, and error handling are exercised
without needing running services.
"""

import importlib.util
import json
import pathlib
import sys
import warnings
from unittest.mock import MagicMock

# Suppress deprecation warnings from third-party libraries (e.g. huggingface_hub)
# that are triggered during module import but are not related to the code under test.
warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub")

import networkx as nx
import pytest

# ---------------------------------------------------------------------------
# Additional sys.modules mocks needed beyond what conftest already provides.
#
# conftest.py (same directory) mocks the heavy packages listed in
# _modules_to_mock.  We need a few more to satisfy index.py and
# task_executor.py's import-time dependencies.
# ---------------------------------------------------------------------------
_EXTRA_MOCKS = [
    # for index.py
    "api.db.services.document_service",
    # for task_executor.py
    "api.db",
    "api.db.services.knowledgebase_service",
    "api.db.services.pipeline_operation_log_service",
    "api.db.joint_services",
    "api.db.joint_services.memory_message_service",
    "api.db.joint_services.tenant_model_service",
    "api.db.services.doc_metadata_service",
    "api.db.services.llm_service",
    "api.db.services.file2document_service",
    "api.db.db_models",
    "common.metadata_utils",
    "common.log_utils",
    "common.config_utils",
    "common.versions",
    "common.token_utils",
    "common.signal_utils",
    "common.exceptions",
    "common.constants",
    "rag.utils.base64_image",

    "rag.prompts.generator",
    "rag.raptor",
    "rag.app",
    "rag.graphrag.utils",
]
for _m in _EXTRA_MOCKS:
    if _m not in sys.modules:
        sys.modules[_m] = MagicMock()

# ---------------------------------------------------------------------------
# Load the real implementations via importlib.
# ---------------------------------------------------------------------------
_ROOT = pathlib.Path(__file__).parents[4]


def _load_module(dotted_name: str, rel_path: str):
    path = _ROOT / rel_path
    spec = importlib.util.spec_from_file_location(dotted_name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[dotted_name] = mod
    spec.loader.exec_module(mod)
    return mod


_index_mod = _load_module("rag.graphrag.general.index", "rag/graphrag/general/index.py")
_executor_mod = _load_module("rag.svr.task_executor", "rag/svr/task_executor.py")

load_subgraph_from_store = _index_mod.load_subgraph_from_store
has_raptor_chunks = _executor_mod.has_raptor_chunks

# settings is a MagicMock installed by conftest; grab it to monkeypatch docStoreConn.
import common.settings as _settings  # noqa: E402

# Ensure docStoreConn is a MagicMock so monkeypatch.setattr works in all environments.
if not isinstance(_settings.docStoreConn, MagicMock):
    _settings.docStoreConn = MagicMock()


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _make_subgraph(doc_id: str) -> nx.Graph:
    sg = nx.Graph()
    sg.add_node("ENTITY_A", description="test entity A", source_id=[doc_id])
    sg.add_node("ENTITY_B", description="test entity B", source_id=[doc_id])
    sg.add_edge("ENTITY_A", "ENTITY_B", description="related", source_id=[doc_id], weight=1.0, keywords=[])
    sg.graph["source_id"] = [doc_id]
    return sg


def _to_store_content(sg: nx.Graph) -> str:
    return json.dumps(nx.node_link_data(sg, edges="edges"), ensure_ascii=False)


def _single_page_mocks(field_map: dict):
    """search + get_fields mocks that simulate a single-page result."""
    sentinel = object()
    call_count = {"n": 0}

    def _get_fields(_res, _fields):
        call_count["n"] += 1
        return field_map if call_count["n"] == 1 else {}

    return MagicMock(return_value=sentinel), MagicMock(side_effect=_get_fields)


# ---------------------------------------------------------------------------
# Tests for load_subgraph_from_store  (rag/graphrag/general/index.py)
# ---------------------------------------------------------------------------

class TestLoadSubgraphFromStore:

    @pytest.mark.p1
    @pytest.mark.asyncio
    async def test_loads_existing_subgraph(self, monkeypatch):
        """Subgraph present in the store is returned as nx.Graph."""
        doc_id = "doc_001"
        sg = _make_subgraph(doc_id)
        field_map = {"chunk_001": {"content_with_weight": _to_store_content(sg), "source_id": [doc_id]}}
        s, gf = _single_page_mocks(field_map)
        monkeypatch.setattr(_settings.docStoreConn, "search", s)
        monkeypatch.setattr(_settings.docStoreConn, "get_fields", gf)

        result = await load_subgraph_from_store("t1", "kb1", doc_id)

        assert result is not None and isinstance(result, nx.Graph)
        assert result.has_node("ENTITY_A") and result.has_node("ENTITY_B")
        assert result.graph["source_id"] == [doc_id]

    @pytest.mark.p1
    @pytest.mark.asyncio
    async def test_returns_none_when_no_subgraph(self, monkeypatch):
        """Empty store returns None without raising."""
        s, gf = _single_page_mocks({})
        monkeypatch.setattr(_settings.docStoreConn, "search", s)
        monkeypatch.setattr(_settings.docStoreConn, "get_fields", gf)

        assert await load_subgraph_from_store("t1", "kb1", "doc_missing") is None

    @pytest.mark.p2
    @pytest.mark.asyncio
    async def test_passes_doc_id_in_search_condition(self, monkeypatch):
        """source_id (== doc_id) is included in the search condition so the doc
        store filters results directly rather than fetching all subgraphs."""
        captured = {}

        def _capture(fields, filters, condition, *_a, **_kw):
            captured["condition"] = condition
            return object()

        sg = _make_subgraph("doc_b")
        monkeypatch.setattr(_settings.docStoreConn, "search", _capture)
        monkeypatch.setattr(_settings.docStoreConn, "get_fields",
                            MagicMock(return_value={"chunk_b": {"content_with_weight": _to_store_content(sg), "source_id": ["doc_b"]}}))

        result = await load_subgraph_from_store("t1", "kb1", "doc_b")
        assert result is not None and result.graph["source_id"] == ["doc_b"]
        assert captured["condition"]["source_id"] == ["doc_b"]

    @pytest.mark.p2
    @pytest.mark.asyncio
    async def test_skips_malformed_json_returns_none(self, monkeypatch):
        """Malformed JSON is logged and skipped; None is returned (not raised)."""
        field_map = {"chunk_bad": {"content_with_weight": "not valid json{{{", "source_id": ["doc_bad"]}}
        s, gf = _single_page_mocks(field_map)
        monkeypatch.setattr(_settings.docStoreConn, "search", s)
        monkeypatch.setattr(_settings.docStoreConn, "get_fields", gf)

        assert await load_subgraph_from_store("t1", "kb1", "doc_bad") is None

    @pytest.mark.p2
    @pytest.mark.asyncio
    async def test_issues_single_query_with_limit_one(self, monkeypatch):
        """Exactly one search call is issued with limit=1 — the doc store index
        does the filtering, so no pagination is required."""
        doc_id = "doc_single"
        sg = _make_subgraph(doc_id)
        search_calls: list[tuple] = []

        def _search(fields, filters, condition, order, orderby, offset, limit, *_a, **_kw):
            search_calls.append((offset, limit))
            return object()

        monkeypatch.setattr(_settings.docStoreConn, "search", _search)
        monkeypatch.setattr(_settings.docStoreConn, "get_fields",
                            MagicMock(return_value={"chunk_t": {"content_with_weight": _to_store_content(sg), "source_id": [doc_id]}}))

        result = await load_subgraph_from_store("t1", "kb1", doc_id)
        assert result is not None
        assert len(search_calls) == 1, "must issue exactly one query"
        assert search_calls[0] == (0, 1), "must use offset=0, limit=1"

    @pytest.mark.p2
    @pytest.mark.asyncio
    async def test_doc_store_exception_returns_none(self, monkeypatch):
        """A doc-store exception is caught; None is returned safely."""
        monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(side_effect=RuntimeError("db down")))
        assert await load_subgraph_from_store("t1", "kb1", "doc_001") is None


# ---------------------------------------------------------------------------
# Tests for has_raptor_chunks  (rag/svr/task_executor.py)
# ---------------------------------------------------------------------------

class TestHasRaptorChunks:

    @pytest.mark.p1
    @pytest.mark.asyncio
    async def test_returns_true_when_raptor_chunk_exists(self, monkeypatch):
        """Doc store returns a RAPTOR row -> True."""
        monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(return_value=object()))
        monkeypatch.setattr(_settings.docStoreConn, "get_fields",
                            MagicMock(return_value={"chunk_r": {"raptor_kwd": "raptor"}}))

        assert await has_raptor_chunks("doc_001", "t1", "kb1") is True

    @pytest.mark.p1
    @pytest.mark.asyncio
    async def test_returns_false_when_no_raptor_chunks(self, monkeypatch):
        """Doc store returns empty -> False."""
        monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(return_value=object()))
        monkeypatch.setattr(_settings.docStoreConn, "get_fields", MagicMock(return_value={}))

        assert await has_raptor_chunks("doc_001", "t1", "kb1") is False

    @pytest.mark.p1
    @pytest.mark.asyncio
    async def test_queries_specifically_for_raptor_kwd(self, monkeypatch):
        """raptor_kwd is in the search condition so non-RAPTOR leading chunks
        cannot produce a false-negative."""
        captured = {}

        def _capture(fields, filters, condition, *_a, **_kw):
            captured["condition"] = condition
            return object()

        monkeypatch.setattr(_settings.docStoreConn, "search", _capture)
        monkeypatch.setattr(_settings.docStoreConn, "get_fields", MagicMock(return_value={}))

        await has_raptor_chunks("doc_001", "t1", "kb1")
        assert captured["condition"] == {"doc_id": "doc_001", "raptor_kwd": ["raptor"]}

    @pytest.mark.p2
    @pytest.mark.asyncio
    async def test_returns_false_on_doc_store_exception(self, monkeypatch):
        """Exception is caught; False is returned without crashing."""
        monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(side_effect=RuntimeError("db down")))
        assert await has_raptor_chunks("doc_001", "t1", "kb1") is False


# ---------------------------------------------------------------------------
# End-to-end workflow test
# ---------------------------------------------------------------------------

class TestCheckpointResumeWorkflow:

    @pytest.mark.p1
    @pytest.mark.asyncio
    async def test_resume_finds_completed_docs_skips_new_ones(self, monkeypatch):
        """3 docs completed before crash; on resume each is found, new doc is not."""
        completed = ["doc_1", "doc_2", "doc_3"]
        field_map = {
            f"chunk_{d}": {"content_with_weight": _to_store_content(_make_subgraph(d)), "source_id": [d]}
            for d in completed
        }
        # The doc store filters by source_id (doc_id) directly, so get_fields
        # should return only the matching chunk for each call.
        def _get_fields_by_doc(res, fields):
            # res is the sentinel from search; extract the doc_id it was called with
            return {k: v for k, v in field_map.items() if v["source_id"] == [_get_fields_by_doc.last_doc_id]}

        def _search(fields, filters, condition, *_a, **_kw):
            _get_fields_by_doc.last_doc_id = (condition or {}).get("source_id", [None])[0]
            return object()

        monkeypatch.setattr(_settings.docStoreConn, "search", _search)
        monkeypatch.setattr(_settings.docStoreConn, "get_fields", _get_fields_by_doc)

        for doc_id in completed:
            result = await load_subgraph_from_store("t1", "kb1", doc_id)
            assert result is not None and result.graph["source_id"] == [doc_id]

        assert await load_subgraph_from_store("t1", "kb1", "doc_4_new") is None