hkuds--lightrag/tests/tools/test_rebuild_vdb.py

"""Tests for the offline VDB rebuild tool and the fail-loud merge semantics.

Covers:
- rebuild: payload fidelity against the authoritative write points,
  bidirectional edge dedup, AGE quote stripping, dirty-data skipping,
  drop-before-upsert ordering, batching with periodic flushes, and
  error collection on persistently failing batches.
- check: consistent stores, missing-record detection, legacy reverse
  relation ids not being misreported, and batching.
- merge: _merge_entities_impl raising VectorStorageConsistencyError on
  persistent VDB failure without deleting source entities.
"""

import pytest
from unittest.mock import AsyncMock, MagicMock

import lightrag.tools.rebuild_vdb as rebuild_vdb
from lightrag.tools.rebuild_vdb import (
    _strip_agtype_quotes,
    check_vdb_consistency,
    rebuild_chunks_vdb,
    rebuild_entities_vdb,
    rebuild_relationships_vdb,
)
from lightrag.utils import (
    VectorStorageConsistencyError,
    compute_mdhash_id,
    make_relation_vdb_ids,
)

pytestmark = pytest.mark.offline


# ---------------------------------------------------------------------------
# Mocks
# ---------------------------------------------------------------------------


class MockVDB:
    """Vector storage mock: records upserts, aligned-None get_by_ids."""

    def __init__(self):
        self.global_config = {}  # no tokenizer -> _truncate_vdb_content is a no-op
        self.records = {}
        self.call_order = []

        async def _drop():
            self.call_order.append("drop")
            self.records.clear()
            return {"status": "success", "message": "data dropped"}

        async def _upsert(payload):
            self.call_order.append("upsert")
            self.records.update(payload)

        async def _get_by_ids(ids):
            return [self.records.get(i) for i in ids]

        self.drop = AsyncMock(side_effect=_drop)
        self.upsert = AsyncMock(side_effect=_upsert)
        self.get_by_ids = AsyncMock(side_effect=_get_by_ids)
        self.delete = AsyncMock()
        self.index_done_callback = AsyncMock()


def make_graph(nodes=None, edges=None):
    graph = MagicMock()
    graph.get_all_nodes = AsyncMock(return_value=nodes or [])
    graph.get_all_edges = AsyncMock(return_value=edges or [])
    return graph


class _FakeLock:
    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        return False


class JsonKVStorage:
    """Minimal stand-in; the class NAME drives enumerate_kv_keys dispatch."""

    def __init__(self, data):
        self._data = data
        self._storage_lock = _FakeLock()

    async def get_by_ids(self, ids):
        return [self._data.get(i) for i in ids]


def node(name, **overrides):
    data = {
        "entity_id": name,
        "description": f"description of {name}",
        "entity_type": "person",
        "source_id": "chunk-abc",
        "file_path": "doc.txt",
    }
    data.update(overrides)
    return data


def edge(src, tgt, **overrides):
    data = {
        "source": src,
        "target": tgt,
        "description": f"{src} knows {tgt}",
        "keywords": "knows",
        "source_id": "chunk-abc",
        "weight": 2.0,
        "file_path": "doc.txt",
    }
    data.update(overrides)
    return data


# ---------------------------------------------------------------------------
# Rebuild: entities
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_rebuild_entities_payload_matches_authoritative_format():
    graph = make_graph(nodes=[node("Alice")])
    vdb = MockVDB()

    stats = await rebuild_entities_vdb(graph, vdb, {})

    entity_id = compute_mdhash_id("Alice", prefix="ent-")
    assert stats["rebuilt"] == 1
    assert stats["skipped"] == 0
    assert vdb.records == {
        entity_id: {
            "entity_name": "Alice",
            "entity_type": "person",
            "content": "Alice\ndescription of Alice",
            "source_id": "chunk-abc",
            "description": "description of Alice",
            "file_path": "doc.txt",
        }
    }


@pytest.mark.asyncio
async def test_rebuild_entities_handles_missing_optional_fields():
    graph = make_graph(
        nodes=[{"id": "Bob"}]  # only "id", no entity_id, no optional fields
    )
    vdb = MockVDB()

    await rebuild_entities_vdb(graph, vdb, {})

    record = vdb.records[compute_mdhash_id("Bob", prefix="ent-")]
    assert record["entity_name"] == "Bob"
    assert record["content"] == "Bob\n"
    assert record["entity_type"] == ""
    assert record["source_id"] == ""
    assert record["file_path"] == ""


@pytest.mark.asyncio
async def test_rebuild_entities_skips_dirty_nodes():
    graph = make_graph(nodes=[node("Alice"), {"description": "no id"}, node("  ")])
    vdb = MockVDB()

    stats = await rebuild_entities_vdb(graph, vdb, {})

    assert stats["source_total"] == 3
    assert stats["rebuilt"] == 1
    assert stats["skipped"] == 2
    assert list(vdb.records) == [compute_mdhash_id("Alice", prefix="ent-")]


@pytest.mark.asyncio
async def test_rebuild_drops_once_before_any_upsert():
    graph = make_graph(nodes=[node("Alice"), node("Bob")])
    vdb = MockVDB()

    await rebuild_entities_vdb(graph, vdb, {}, batch_size=1)

    assert vdb.drop.await_count == 1
    assert vdb.call_order[0] == "drop"
    assert vdb.call_order.count("upsert") == 2


@pytest.mark.asyncio
async def test_rebuild_batches_and_periodic_flush():
    graph = make_graph(nodes=[node(f"E{i:03d}") for i in range(25)])
    vdb = MockVDB()

    stats = await rebuild_entities_vdb(graph, vdb, {}, batch_size=1)

    assert stats["batches"] == 25
    assert vdb.upsert.await_count == 25
    # periodic flush at batches 10 and 20, plus the final flush
    assert vdb.index_done_callback.await_count == 3


@pytest.mark.asyncio
async def test_rebuild_collects_batch_errors_and_continues(monkeypatch):
    # Single-attempt wrapper: no retry delays in tests
    async def _single_attempt(operation, **kwargs):
        await operation()

    monkeypatch.setattr(
        rebuild_vdb, "safe_vdb_operation_with_exception", _single_attempt
    )

    graph = make_graph(nodes=[node("Alice"), node("Bob"), node("Carol")])
    vdb = MockVDB()
    poisoned_id = compute_mdhash_id("Bob", prefix="ent-")

    async def _failing_upsert(payload):
        if poisoned_id in payload:
            raise RuntimeError("embedder down")
        vdb.records.update(payload)

    vdb.upsert = AsyncMock(side_effect=_failing_upsert)

    stats = await rebuild_entities_vdb(graph, vdb, {}, batch_size=1)

    assert stats["rebuilt"] == 2
    assert stats["failed_batches"] == 1
    assert len(stats["errors"]) == 1
    assert stats["errors"][0]["error_msg"] == "embedder down"
    assert poisoned_id not in vdb.records
    assert len(vdb.records) == 2


# ---------------------------------------------------------------------------
# Rebuild: relationships
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_rebuild_relationships_payload_normalized():
    # Endpoints arrive in reverse lexicographic order; payload must be sorted
    graph = make_graph(edges=[edge("Bob", "Alice")])
    vdb = MockVDB()

    stats = await rebuild_relationships_vdb(graph, vdb, {})

    rel_id = compute_mdhash_id("Alice" + "Bob", prefix="rel-")
    assert stats["rebuilt"] == 1
    assert vdb.records == {
        rel_id: {
            "src_id": "Alice",
            "tgt_id": "Bob",
            "source_id": "chunk-abc",
            "content": "knows\tAlice\nBob\nBob knows Alice",
            "keywords": "knows",
            "description": "Bob knows Alice",
            "weight": 2.0,
            "file_path": "doc.txt",
        }
    }


@pytest.mark.asyncio
async def test_rebuild_relationships_weight_fallback():
    graph = make_graph(edges=[edge("A", "B", weight="3"), edge("C", "D", weight=None)])
    vdb = MockVDB()

    await rebuild_relationships_vdb(graph, vdb, {})

    assert vdb.records[compute_mdhash_id("AB", prefix="rel-")]["weight"] == 3.0
    assert vdb.records[compute_mdhash_id("CD", prefix="rel-")]["weight"] == 1.0


@pytest.mark.asyncio
async def test_rebuild_relationships_dedupes_bidirectional_edges():
    # Neo4j/Memgraph return each undirected edge once per direction
    graph = make_graph(edges=[edge("Alice", "Bob"), edge("Bob", "Alice")])
    vdb = MockVDB()

    stats = await rebuild_relationships_vdb(graph, vdb, {})

    assert stats["source_total"] == 2
    assert stats["rebuilt"] == 1
    assert stats["duplicates"] == 1


@pytest.mark.asyncio
async def test_rebuild_relationships_strips_age_quotes():
    # PG/AGE agtype::text casts wrap endpoint strings in double quotes
    graph = make_graph(edges=[edge('"Alice"', '"Bob"')])
    vdb = MockVDB()

    await rebuild_relationships_vdb(graph, vdb, {})

    rel_id = compute_mdhash_id("AliceBob", prefix="rel-")
    assert vdb.records[rel_id]["src_id"] == "Alice"
    assert vdb.records[rel_id]["tgt_id"] == "Bob"


@pytest.mark.asyncio
async def test_rebuild_relationships_skips_edges_without_endpoints():
    graph = make_graph(edges=[edge("Alice", "Bob"), {"description": "no endpoints"}])
    vdb = MockVDB()

    stats = await rebuild_relationships_vdb(graph, vdb, {})

    assert stats["rebuilt"] == 1
    assert stats["skipped"] == 1


def test_strip_agtype_quotes():
    assert _strip_agtype_quotes('"Alice"') == "Alice"
    assert _strip_agtype_quotes("Alice") == "Alice"
    assert _strip_agtype_quotes('"') == '"'
    assert _strip_agtype_quotes(None) is None
    assert _strip_agtype_quotes(3) == 3


# ---------------------------------------------------------------------------
# Rebuild: chunks
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_rebuild_chunks_payload_from_kv_records():
    kv = JsonKVStorage(
        {
            "chunk-1": {
                "content": "chunk text",
                "full_doc_id": "doc-1",
                "file_path": "doc.txt",
                "tokens": 3,
                "chunk_order_index": 0,
            },
        }
    )
    vdb = MockVDB()

    stats = await rebuild_chunks_vdb(kv, vdb)

    assert stats["source_total"] == 1
    assert stats["rebuilt"] == 1
    record = vdb.records["chunk-1"]
    assert record["content"] == "chunk text"
    assert record["full_doc_id"] == "doc-1"
    assert record["file_path"] == "doc.txt"
    # extra pipeline fields pass through, as in the authoritative upsert
    assert record["tokens"] == 3
    assert record["chunk_order_index"] == 0
    assert vdb.drop.await_count == 1


@pytest.mark.asyncio
async def test_rebuild_chunks_skips_records_without_content():
    kv = JsonKVStorage(
        {
            "chunk-1": {"content": "ok", "full_doc_id": "doc-1"},
            "chunk-2": {"full_doc_id": "doc-1"},  # no content
        }
    )
    vdb = MockVDB()

    stats = await rebuild_chunks_vdb(kv, vdb)

    assert stats["rebuilt"] == 1
    assert stats["skipped"] == 1
    assert "chunk-2" not in vdb.records


@pytest.mark.asyncio
async def test_rebuild_chunks_covers_all_id_schemes():
    # Chunks live under several id schemes that no single prefix matches:
    # custom KG ("chunk-<hash>"), the text pipeline ("{doc_id}-chunk-{order}",
    # which does NOT start with "chunk-"), and multimodal
    # ("{doc_id}-mm-<modality>-{order}"). A prefix-only filter dropped the
    # pipeline/multimodal chunks and rebuilt to zero. All schemes must rebuild.
    kv = JsonKVStorage(
        {
            "chunk-deadbeef": {"content": "custom kg", "full_doc_id": "doc-xyz"},
            "doc-abc123-chunk-000": {"content": "first", "full_doc_id": "doc-abc123"},
            "doc-abc123-chunk-001": {"content": "second", "full_doc_id": "doc-abc123"},
            "doc-abc123-mm-drawing-000": {
                "content": "image caption",
                "full_doc_id": "doc-abc123",
            },
        }
    )
    vdb = MockVDB()

    stats = await rebuild_chunks_vdb(kv, vdb)

    assert stats["source_total"] == 4
    assert stats["rebuilt"] == 4
    assert set(vdb.records) == {
        "chunk-deadbeef",
        "doc-abc123-chunk-000",
        "doc-abc123-chunk-001",
        "doc-abc123-mm-drawing-000",
    }


# ---------------------------------------------------------------------------
# Consistency check
# ---------------------------------------------------------------------------


def seeded_vdbs(nodes, edges):
    """Build entity/relation VDB mocks already containing all graph records."""
    entities_vdb = MockVDB()
    for n in nodes:
        entities_vdb.records[compute_mdhash_id(n["entity_id"], prefix="ent-")] = {
            "entity_name": n["entity_id"]
        }
    relationships_vdb = MockVDB()
    for e in edges:
        src, tgt = sorted([e["source"], e["target"]])
        relationships_vdb.records[compute_mdhash_id(src + tgt, prefix="rel-")] = {
            "src_id": src,
            "tgt_id": tgt,
        }
    return entities_vdb, relationships_vdb


@pytest.mark.asyncio
async def test_check_all_consistent():
    nodes = [node("Alice"), node("Bob")]
    edges = [edge("Alice", "Bob")]
    graph = make_graph(nodes=nodes, edges=edges)
    entities_vdb, relationships_vdb = seeded_vdbs(nodes, edges)

    report = await check_vdb_consistency(graph, entities_vdb, relationships_vdb)

    assert report["consistent"] is True
    assert report["graph_entities"] == 2
    assert report["graph_relations"] == 1
    assert report["missing_entities"] == 0
    assert report["missing_relations"] == 0


@pytest.mark.asyncio
async def test_check_detects_missing_records():
    nodes = [node("Alice"), node("Bob")]
    edges = [edge("Alice", "Bob"), edge("Bob", "Carol")]
    graph = make_graph(nodes=nodes, edges=edges)
    entities_vdb, relationships_vdb = seeded_vdbs(nodes[:1], edges[:1])

    report = await check_vdb_consistency(graph, entities_vdb, relationships_vdb)

    assert report["consistent"] is False
    assert report["missing_entities"] == 1
    assert report["missing_entity_names"] == ["Bob"]
    assert report["missing_relations"] == 1
    assert report["missing_relation_pairs"] == ["Bob ~ Carol"]


@pytest.mark.asyncio
async def test_check_accepts_legacy_reverse_relation_id():
    # Legacy custom-KG imports hashed the relation in original endpoint
    # order; the VDB holds only the reverse-order id. Not an inconsistency.
    graph = make_graph(edges=[edge("Bob", "Alice")])
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    reverse_id = make_relation_vdb_ids("Bob", "Alice")[1]
    relationships_vdb.records[reverse_id] = {"src_id": "Bob", "tgt_id": "Alice"}

    report = await check_vdb_consistency(graph, entities_vdb, relationships_vdb)

    assert report["missing_relations"] == 0
    assert report["consistent"] is True


@pytest.mark.asyncio
async def test_check_batches_requests():
    nodes = [node(f"E{i:03d}") for i in range(5)]
    graph = make_graph(nodes=nodes)
    entities_vdb, relationships_vdb = seeded_vdbs(nodes, [])

    report = await check_vdb_consistency(
        graph, entities_vdb, relationships_vdb, batch_size=2
    )

    assert report["consistent"] is True
    assert entities_vdb.get_by_ids.await_count == 3  # ceil(5 / 2)
    assert all(
        len(call.args[0]) <= 2 for call in entities_vdb.get_by_ids.await_args_list
    )


# ---------------------------------------------------------------------------
# Merge: fail-loud semantics in _merge_entities_impl
# ---------------------------------------------------------------------------


def make_merge_graph():
    """Graph mock for merging source 'Bob' into new target 'Alice'.

    Bob has one edge to Carol; Alice does not exist yet.
    """
    graph = MagicMock()
    graph.has_node = AsyncMock(side_effect=lambda name: name == "Bob")
    graph.get_node = AsyncMock(
        return_value={
            "entity_id": "Bob",
            "description": "desc",
            "entity_type": "person",
            "source_id": "chunk-abc",
            "file_path": "doc.txt",
        }
    )
    graph.get_node_edges = AsyncMock(return_value=[("Bob", "Carol")])
    graph.get_edge = AsyncMock(
        return_value={
            "description": "knows",
            "keywords": "kw",
            "source_id": "chunk-abc",
            "weight": 1.0,
            "file_path": "doc.txt",
        }
    )
    graph.upsert_node = AsyncMock()
    graph.upsert_edge = AsyncMock()
    graph.delete_node = AsyncMock()
    graph.index_done_callback = AsyncMock()
    return graph


@pytest.fixture
def single_attempt_vdb_ops(monkeypatch):
    """Replace the retry wrapper with a single attempt to keep tests fast."""
    import lightrag.utils_graph as utils_graph

    async def _single_attempt(operation, **kwargs):
        await operation()

    monkeypatch.setattr(
        utils_graph, "safe_vdb_operation_with_exception", _single_attempt
    )


@pytest.mark.asyncio
async def test_merge_relation_vdb_delete_failure_raises_consistency_error(
    single_attempt_vdb_ops,
):
    # The step-7 delete of stale relation records runs after the graph is
    # updated; a failure there must fail loud (and not be swallowed upstream).
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    relationships_vdb.delete = AsyncMock(side_effect=RuntimeError("vdb delete down"))

    with pytest.raises(VectorStorageConsistencyError) as excinfo:
        await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )

    assert "lightrag-rebuild-vdb" in str(excinfo.value)
    relationships_vdb.delete.assert_awaited()
    # Failure is before source-entity deletion (step 10)
    graph.delete_node.assert_not_awaited()
    entities_vdb.delete.assert_not_awaited()


@pytest.mark.asyncio
async def test_merge_relation_vdb_failure_raises_consistency_error(
    single_attempt_vdb_ops,
):
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    relationships_vdb.upsert = AsyncMock(side_effect=RuntimeError("embedder down"))

    with pytest.raises(VectorStorageConsistencyError) as excinfo:
        await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )

    # Fail-loud guidance in the message
    assert "lightrag-rebuild-vdb" in str(excinfo.value)
    # Graph state was written and is kept (no rollback)
    graph.upsert_edge.assert_awaited()
    # Source entities were NOT deleted (step 10 never reached)
    graph.delete_node.assert_not_awaited()
    entities_vdb.delete.assert_not_awaited()


@pytest.mark.asyncio
async def test_merge_entity_vdb_failure_raises_consistency_error(
    single_attempt_vdb_ops,
):
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    entities_vdb.upsert = AsyncMock(side_effect=RuntimeError("embedder down"))

    with pytest.raises(VectorStorageConsistencyError) as excinfo:
        await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )

    assert "lightrag-rebuild-vdb" in str(excinfo.value)
    # Relation VDB write succeeded before the entity failure
    relationships_vdb.upsert.assert_awaited()
    graph.delete_node.assert_not_awaited()
    entities_vdb.delete.assert_not_awaited()


@pytest.mark.asyncio
async def test_merge_source_entity_delete_failure_raises_consistency_error(
    single_attempt_vdb_ops,
):
    # Step 10 deletes the source node from the graph first, then its vector
    # record. A delete failure here happens AFTER the graph node is gone, so it
    # must fail loud with a message that does not claim the source still exists.
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    entities_vdb.delete = AsyncMock(side_effect=RuntimeError("vdb delete down"))

    with pytest.raises(VectorStorageConsistencyError) as excinfo:
        await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )

    message = str(excinfo.value)
    assert "lightrag-rebuild-vdb" in message
    # The source node was already removed; the message must not claim otherwise.
    graph.delete_node.assert_awaited_with("Bob")
    assert "were not deleted" not in message
    entities_vdb.delete.assert_awaited()


@pytest.mark.asyncio
async def test_merge_final_persist_failure_raises_consistency_error(
    single_attempt_vdb_ops,
):
    # The step-11 _persist_graph_updates flush runs after source deletion; a
    # failure there must also fail loud (not a raw exception swallowed upstream).
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    # Succeed at the 8b flush, fail at the final persist flush (step 11).
    entities_vdb.index_done_callback = AsyncMock(
        side_effect=[None, RuntimeError("flush down")]
    )

    with pytest.raises(VectorStorageConsistencyError) as excinfo:
        await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )

    message = str(excinfo.value)
    assert "lightrag-rebuild-vdb" in message
    # Source deletion already happened before the final persist.
    graph.delete_node.assert_awaited_with("Bob")
    assert "were not deleted" not in message


@pytest.mark.asyncio
async def test_merge_success_path_unaffected(single_attempt_vdb_ops):
    import lightrag.utils_graph as utils_graph
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()

    async def _noop_persist(**kwargs):
        return None

    async def _entity_info(*args, **kwargs):
        return {"entity_name": "Alice"}

    # Avoid index_done_callback/inspection helpers needing full storages
    orig_persist = utils_graph._persist_graph_updates
    orig_info = utils_graph.get_entity_info
    utils_graph._persist_graph_updates = _noop_persist
    utils_graph.get_entity_info = _entity_info
    try:
        result = await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )
    finally:
        utils_graph._persist_graph_updates = orig_persist
        utils_graph.get_entity_info = orig_info

    assert result == {"entity_name": "Alice"}
    # Both VDB writes happened and the source entity was deleted
    relationships_vdb.upsert.assert_awaited()
    entities_vdb.upsert.assert_awaited()
    graph.delete_node.assert_awaited_with("Bob")


@pytest.mark.asyncio
async def test_merge_deferred_flush_failure_raises_before_delete(
    single_attempt_vdb_ops,
):
    # Deferred-embedding backends (nano/faiss) buffer in upsert() and only embed
    # in index_done_callback, so an embedder outage surfaces at flush time. The
    # fail-loud guarantee must still hold: raise before deleting source entities.
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    # upsert succeeds (buffers); the embedding failure happens at flush.
    relationships_vdb.index_done_callback = AsyncMock(
        side_effect=RuntimeError("embedder down")
    )

    with pytest.raises(VectorStorageConsistencyError) as excinfo:
        await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )

    assert "lightrag-rebuild-vdb" in str(excinfo.value)
    # upsert succeeded; the failure is the deferred flush
    relationships_vdb.upsert.assert_awaited()
    relationships_vdb.index_done_callback.assert_awaited()
    # Source entities NOT deleted (step 10 never reached), so the message holds
    graph.delete_node.assert_not_awaited()
    entities_vdb.delete.assert_not_awaited()


@pytest.mark.asyncio
async def test_merge_entity_deferred_flush_failure_raises_before_delete(
    single_attempt_vdb_ops,
):
    from lightrag.utils_graph import _merge_entities_impl

    graph = make_merge_graph()
    entities_vdb = MockVDB()
    relationships_vdb = MockVDB()
    entities_vdb.index_done_callback = AsyncMock(
        side_effect=RuntimeError("embedder down")
    )

    with pytest.raises(VectorStorageConsistencyError) as excinfo:
        await _merge_entities_impl(
            graph, entities_vdb, relationships_vdb, ["Bob"], "Alice"
        )

    assert "lightrag-rebuild-vdb" in str(excinfo.value)
    # Relation flush succeeded before the entity flush failed
    relationships_vdb.index_done_callback.assert_awaited()
    entities_vdb.index_done_callback.assert_awaited()
    graph.delete_node.assert_not_awaited()
    entities_vdb.delete.assert_not_awaited()


# ---------------------------------------------------------------------------
# Setup: check-only (no api extra) stub embedding function
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_check_only_stub_carries_embedding_model_name(monkeypatch):
    # When the api extra is unavailable the tool runs check-only with a stub
    # embedding function. Qdrant/PostgreSQL locate the collection/table from
    # model_name + embedding_dim, so the stub must carry the configured
    # EMBEDDING_MODEL — otherwise check-only probes the legacy-named collection
    # and misreports records as missing.
    import lightrag.kg.factory as kg_factory
    from lightrag.tools.rebuild_vdb import RebuildTool

    monkeypatch.setenv("EMBEDDING_MODEL", "text-embedding-3-large")
    monkeypatch.setenv("EMBEDDING_DIM", "3072")

    captured_funcs = []

    class _DummyStorage:
        def __init__(self, *, embedding_func, **kwargs):
            captured_funcs.append(embedding_func)

        async def initialize(self):
            return None

    monkeypatch.setattr(kg_factory, "get_storage_class", lambda name: _DummyStorage)

    tool = RebuildTool()
    # Simulate the api extra being unavailable (check-only mode).
    monkeypatch.setattr(tool, "build_embedding_func", lambda: None)

    ok = await tool.setup_storages()

    assert ok is True
    assert tool.embedding_available is False
    assert tool.embedding_func.model_name == "text-embedding-3-large"
    assert tool.embedding_func.embedding_dim == 3072
    # Every storage received the stub carrying the configured model name.
    assert captured_funcs
    assert all(f.model_name == "text-embedding-3-large" for f in captured_funcs)


@pytest.mark.asyncio
async def test_check_only_stub_model_name_none_when_unset(monkeypatch):
    # EMBEDDING_MODEL unset -> model_name None (provider default), matching the
    # server factory; "None" literal is also treated as unset.
    import lightrag.kg.factory as kg_factory
    from lightrag.tools.rebuild_vdb import RebuildTool

    monkeypatch.delenv("EMBEDDING_MODEL", raising=False)

    class _DummyStorage:
        def __init__(self, **kwargs):
            pass

        async def initialize(self):
            return None

    monkeypatch.setattr(kg_factory, "get_storage_class", lambda name: _DummyStorage)

    tool = RebuildTool()
    monkeypatch.setattr(tool, "build_embedding_func", lambda: None)

    await tool.setup_storages()

    assert tool.embedding_func.model_name is None


# ---------------------------------------------------------------------------
# CLI exit-code semantics: run() must report success/failure
# ---------------------------------------------------------------------------


def _runnable_tool(monkeypatch, inputs):
    """A RebuildTool with shared-storage init/finalize and prompts stubbed."""
    import lightrag.kg.shared_storage as shared

    monkeypatch.setattr(shared, "initialize_share_data", lambda workers=1: None)
    monkeypatch.setattr(shared, "finalize_share_data", lambda: None)
    monkeypatch.setattr("builtins.input", lambda *a: next(inputs))

    tool = rebuild_vdb.RebuildTool()
    monkeypatch.setattr(tool, "print_header", lambda: None)
    monkeypatch.setattr(tool, "confirm_server_stopped", lambda: True)
    return tool


@pytest.mark.asyncio
async def test_run_returns_true_on_user_cancel(monkeypatch):
    tool = _runnable_tool(monkeypatch, iter([]))
    monkeypatch.setattr(tool, "confirm_server_stopped", lambda: False)
    tool.setup_storages = AsyncMock(return_value=True)

    assert await tool.run() is True
    tool.setup_storages.assert_not_awaited()


@pytest.mark.asyncio
async def test_run_returns_false_on_setup_failure(monkeypatch):
    tool = _runnable_tool(monkeypatch, iter([]))
    tool.setup_storages = AsyncMock(return_value=False)

    assert await tool.run() is False


@pytest.mark.asyncio
async def test_run_returns_false_on_unhandled_exception(monkeypatch):
    tool = _runnable_tool(monkeypatch, iter([]))
    tool.setup_storages = AsyncMock(side_effect=RuntimeError("db down"))

    assert await tool.run() is False


@pytest.mark.asyncio
async def test_run_returns_false_when_rebuild_has_errors(monkeypatch):
    tool = _runnable_tool(monkeypatch, iter(["2", "0"]))
    tool.setup_storages = AsyncMock(return_value=True)
    tool.embedding_available = True
    tool.print_source_counts = AsyncMock()
    monkeypatch.setattr(tool, "confirm_rebuild", lambda targets: True)

    failed = rebuild_vdb._new_stats("entities", 3)
    failed["failed_batches"] = 1
    failed["errors"] = [
        {
            "batch": 1,
            "records_lost": 3,
            "error_type": "RuntimeError",
            "error_msg": "embedder down",
        }
    ]
    tool.run_rebuild_entities_relations = AsyncMock(return_value=[failed])

    assert await tool.run() is False


@pytest.mark.asyncio
async def test_run_returns_true_on_clean_rebuild(monkeypatch):
    tool = _runnable_tool(monkeypatch, iter(["2", "0"]))
    tool.setup_storages = AsyncMock(return_value=True)
    tool.embedding_available = True
    tool.print_source_counts = AsyncMock()
    monkeypatch.setattr(tool, "confirm_rebuild", lambda targets: True)

    ok = rebuild_vdb._new_stats("entities", 3)
    ok["rebuilt"] = 3
    tool.run_rebuild_entities_relations = AsyncMock(return_value=[ok])

    assert await tool.run() is True


@pytest.mark.asyncio
async def test_run_check_only_exits_true(monkeypatch):
    tool = _runnable_tool(monkeypatch, iter(["1", "0"]))
    tool.setup_storages = AsyncMock(return_value=True)
    tool.embedding_available = False
    tool.run_check = AsyncMock()

    assert await tool.run() is True
    tool.run_check.assert_awaited()


def test_main_exits_nonzero_on_failure(monkeypatch):
    monkeypatch.setattr(rebuild_vdb, "load_dotenv", lambda **k: None)
    monkeypatch.setattr(rebuild_vdb, "setup_logger", lambda *a, **k: None)

    async def _fail():
        return False

    monkeypatch.setattr(rebuild_vdb, "async_main", _fail)

    with pytest.raises(SystemExit) as excinfo:
        rebuild_vdb.main()
    assert excinfo.value.code == 1


def test_main_exits_zero_on_success(monkeypatch):
    monkeypatch.setattr(rebuild_vdb, "load_dotenv", lambda **k: None)
    monkeypatch.setattr(rebuild_vdb, "setup_logger", lambda *a, **k: None)

    async def _ok():
        return True

    monkeypatch.setattr(rebuild_vdb, "async_main", _ok)

    # No SystemExit on success.
    rebuild_vdb.main()