tests/test_opensearch_storage.py

"""
Unit tests for OpenSearch storage implementations.

All tests use mocks — no running OpenSearch instance required.
Run with: pytest tests/test_opensearch_storage.py -v
"""

import pytest
from contextlib import asynccontextmanager
from unittest.mock import AsyncMock, patch
import numpy as np

pytest.importorskip(
    "opensearchpy",
    reason="opensearchpy is required for OpenSearch storage tests",
)

from opensearchpy.exceptions import NotFoundError, OpenSearchException  # type: ignore
from lightrag.kg.opensearch_impl import (
    OpenSearchKVStorage,
    OpenSearchDocStatusStorage,
    OpenSearchGraphStorage,
    OpenSearchVectorDBStorage,
    ClientManager,
    _build_index_name,
    _resolve_workspace,
    _sanitize_index_name,
)
from lightrag.base import DocStatus, DocProcessingStatus

pytestmark = pytest.mark.offline


# ---------------------------------------------------------------------------
# Mock the shared storage lock so tests don't need full LightRAG init
# ---------------------------------------------------------------------------


@asynccontextmanager
async def _mock_lock():
    yield


def _mock_lock_factory():
    return _mock_lock()


def _missing_index_error() -> NotFoundError:
    return NotFoundError(404, "index_not_found_exception", "no such index")


@pytest.fixture(autouse=True)
def patch_data_init_lock():
    """Patch get_data_init_lock globally so initialize() works without shared storage."""
    with patch(
        "lightrag.kg.opensearch_impl.get_data_init_lock", side_effect=_mock_lock_factory
    ):
        yield


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


class MockEmbeddingFunc:
    """Mock embedding function that returns random vectors."""

    def __init__(self, dim=128):
        self.embedding_dim = dim
        self.max_token_size = 512
        self.model_name = "mock-embed"

    async def __call__(self, texts, **kwargs):
        return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)


@pytest.fixture
def global_config():
    """Standard global config fixture for all storage tests."""
    return {
        "embedding_batch_num": 10,
        "max_graph_nodes": 1000,
        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
    }


@pytest.fixture
def embed_func():
    """Mock embedding function fixture."""
    return MockEmbeddingFunc()


def _make_client():
    """Create a fully-mocked AsyncOpenSearch client with spec validation."""
    from opensearchpy import AsyncOpenSearch

    client = AsyncMock(spec=AsyncOpenSearch)
    # indices sub-client
    client.indices = AsyncMock()
    client.indices.exists = AsyncMock(return_value=False)
    client.indices.create = AsyncMock()
    client.indices.delete = AsyncMock()
    client.indices.refresh = AsyncMock()
    client.indices.get_mapping = AsyncMock(return_value={})
    # transport for PPL
    client.transport = AsyncMock()
    client.transport.perform_request = AsyncMock(
        side_effect=Exception("PPL not available")
    )
    # document operations
    client.exists = AsyncMock(return_value=False)
    client.index = AsyncMock()
    client.delete = AsyncMock()
    client.delete_by_query = AsyncMock()
    client.get = AsyncMock(
        return_value={
            "_id": "doc1",
            "_source": {"content": "hello", "create_time": 0, "update_time": 0},
        }
    )
    client.mget = AsyncMock(
        return_value={
            "docs": [
                {"_id": "id1", "found": True, "_source": {"content": "c1"}},
                {"_id": "id2", "found": True, "_source": {"content": "c2"}},
            ]
        }
    )
    client.count = AsyncMock(return_value={"count": 5})
    client.search = AsyncMock(
        return_value={
            "hits": {"hits": [], "total": {"value": 0}},
            "aggregations": {
                "status_counts": {"buckets": []},
                "src": {"buckets": []},
                "tgt": {"buckets": []},
                "source_degrees": {"buckets": []},
                "target_degrees": {"buckets": []},
            },
        }
    )
    # PIT operations
    client.create_pit = AsyncMock(return_value={"pit_id": "mock_pit_id_123"})
    client.delete_pit = AsyncMock()
    return client


@pytest.fixture
def mock_client():
    """Fully-mocked AsyncOpenSearch client fixture."""
    return _make_client()


# ---------------------------------------------------------------------------
# Helper utilities
# ---------------------------------------------------------------------------


class TestHelpers:
    """Tests for module-level helper functions (_build_index_name, _resolve_workspace, _sanitize_index_name)."""

    def test_build_index_name_with_workspace(self):
        ws, ns, idx = _build_index_name("myws", "text_chunks")
        assert ws == "myws"
        assert ns == "myws_text_chunks"
        assert idx == _sanitize_index_name("myws_text_chunks")

    def test_build_index_name_no_workspace(self):
        ws, ns, idx = _build_index_name("", "chunks")
        assert ws == ""
        assert idx == _sanitize_index_name("chunks")

    def test_resolve_workspace_env_override(self):
        with patch.dict("os.environ", {"OPENSEARCH_WORKSPACE": "forced"}):
            assert _resolve_workspace("original", "ns") == "forced"

    def test_resolve_workspace_fallback(self):
        with patch.dict("os.environ", {}, clear=True):
            assert _resolve_workspace("original", "ns") == "original"

    def test_sanitize_index_name(self):
        assert _sanitize_index_name("Hello_World") == "hello_world"
        assert _sanitize_index_name("-bad") == "x-bad"
        assert _sanitize_index_name("a.b/c") == "a_b_c"


# ---------------------------------------------------------------------------
# ClientManager
# ---------------------------------------------------------------------------


class TestClientManager:
    """Tests for ClientManager singleton pattern and reference counting."""

    @pytest.mark.asyncio
    async def test_singleton_and_refcount(self):
        ClientManager._instances = {"client": None, "ref_count": 0}
        with patch("lightrag.kg.opensearch_impl.AsyncOpenSearch") as mock_cls:
            mock_cls.return_value = AsyncMock()
            c1 = await ClientManager.get_client()
            c2 = await ClientManager.get_client()
            assert c1 is c2
            assert ClientManager._instances["ref_count"] == 2
            await ClientManager.release_client(c1)
            assert ClientManager._instances["ref_count"] == 1
            await ClientManager.release_client(c2)
            assert ClientManager._instances["ref_count"] == 0
            assert ClientManager._instances["client"] is None

    @pytest.mark.asyncio
    async def test_close_called_on_last_release(self):
        ClientManager._instances = {"client": None, "ref_count": 0}
        with patch("lightrag.kg.opensearch_impl.AsyncOpenSearch") as mock_cls:
            inner = AsyncMock()
            mock_cls.return_value = inner
            c = await ClientManager.get_client()
            await ClientManager.release_client(c)
            inner.close.assert_awaited_once()


# ---------------------------------------------------------------------------
# KV Storage
# ---------------------------------------------------------------------------


class TestKVStorage:
    """Tests for OpenSearchKVStorage CRUD operations, timestamps, refresh behavior."""

    def _make(self, global_config, embed_func, workspace="test"):
        return OpenSearchKVStorage(
            namespace="text_chunks",
            global_config=global_config,
            embedding_func=embed_func,
            workspace=workspace,
        )

    @pytest.mark.asyncio
    async def test_index_name(self, global_config, embed_func):
        s = self._make(global_config, embed_func, workspace="proj_a")
        assert s._index_name == "proj_a_text_chunks"

    @pytest.mark.asyncio
    async def test_initialize_creates_index(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            mock_client.indices.exists.assert_awaited_once()
            mock_client.indices.create.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_initialize_skips_existing_index(
        self, global_config, embed_func, mock_client
    ):
        mock_client.indices.exists = AsyncMock(return_value=True)
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            mock_client.indices.create.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_get_by_id(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {
                        "_id": "doc1",
                        "found": True,
                        "_source": {
                            "content": "hello",
                            "create_time": 0,
                            "update_time": 0,
                        },
                    }
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            doc = await s.get_by_id("doc1")
            assert doc is not None
            assert doc["content"] == "hello"
            assert doc["_id"] == "doc1"
            mock_client.mget.assert_awaited_once_with(
                index=s._index_name, body={"ids": ["doc1"]}
            )

    @pytest.mark.asyncio
    async def test_get_by_id_not_found(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={"docs": [{"_id": "missing", "found": False}]}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.get_by_id("missing") is None
            mock_client.get.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_get_by_ids_preserves_order(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            docs = await s.get_by_ids(["id1", "id2"])
            assert docs[0]["content"] == "c1"
            assert docs[1]["content"] == "c2"

    @pytest.mark.asyncio
    async def test_filter_keys(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {"_id": "a", "found": True},
                    {"_id": "b", "found": False},
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.filter_keys({"a", "b"})
            assert result == {"b"}

    @pytest.mark.asyncio
    async def test_upsert_no_per_operation_refresh(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (1, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.upsert({"k1": {"content": "v1"}})
                _, kwargs = mock_bulk.call_args
                assert "refresh" not in kwargs

    @pytest.mark.asyncio
    async def test_upsert_sets_timestamps(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (1, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.upsert({"k1": {"content": "v1"}})
                actions = mock_bulk.call_args[0][1]
                src = actions[0]["_source"]
                assert "create_time" in src
                assert "update_time" in src

    @pytest.mark.asyncio
    async def test_is_empty(self, global_config, embed_func, mock_client):
        mock_client.count = AsyncMock(return_value={"count": 0})
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.is_empty() is True

    @pytest.mark.asyncio
    async def test_delete(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (2, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.delete(["a", "b"])
                actions = mock_bulk.call_args[0][1]
                assert all(a["_op_type"] == "delete" for a in actions)

    @pytest.mark.asyncio
    async def test_drop(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.drop()
            assert result["status"] == "success"
            mock_client.indices.delete.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_drop_error_marks_index_not_ready_and_next_upsert_recreates_index(
        self, global_config, embed_func, mock_client
    ):
        mock_client.indices.delete = AsyncMock(
            side_effect=OpenSearchException("drop failed")
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (1, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                with patch.object(
                    s, "_create_index_if_not_exists", new_callable=AsyncMock
                ) as mock_create:
                    result = await s.drop()
                    assert result["status"] == "error"
                    assert s._index_ready is False
                    await s.upsert({"k1": {"content": "v1"}})
                    mock_create.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_upsert_after_drop_recreates_index(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (1, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                with patch.object(
                    s, "_create_index_if_not_exists", new_callable=AsyncMock
                ) as mock_create:
                    await s.drop()
                    await s.upsert({"k1": {"content": "v1"}})
                    mock_create.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_reads_short_circuit_after_drop(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.drop()

            assert await s.get_by_id("doc1") is None
            assert await s.get_by_ids(["doc1", "doc2"]) == [None, None]
            assert await s.is_empty() is True

            mock_client.mget.assert_not_awaited()
            mock_client.count.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_read_missing_index_demotes_readiness(
        self, global_config, embed_func, mock_client
    ):
        mock_client.mget = AsyncMock(side_effect=_missing_index_error())
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()

            assert await s.get_by_id("doc1") is None
            assert await s.get_by_id("doc1") is None
            assert s._index_ready is False
            assert mock_client.mget.await_count == 1

    @pytest.mark.asyncio
    async def test_iter_raw_docs_uses_pit_and_search_after(
        self, global_config, embed_func, mock_client
    ):
        mock_client.search = AsyncMock(
            side_effect=[
                {
                    "hits": {
                        "hits": [
                            {"_id": "d1", "_source": {"content": "a"}, "sort": [1]},
                            {"_id": "d2", "_source": {"content": "b"}, "sort": [2]},
                        ]
                    }
                },
                {
                    "hits": {
                        "hits": [
                            {"_id": "d3", "_source": {"content": "c"}, "sort": [3]}
                        ]
                    }
                },
            ]
        )

        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()

            batches = [batch async for batch in s._iter_raw_docs(batch_size=2)]

            assert [[doc["_id"] for doc in batch] for batch in batches] == [
                ["d1", "d2"],
                ["d3"],
            ]
            assert (
                "search_after"
                not in mock_client.search.await_args_list[0].kwargs["body"]
            )
            assert mock_client.search.await_args_list[1].kwargs["body"][
                "search_after"
            ] == [2]
            mock_client.create_pit.assert_awaited_once()
            mock_client.delete_pit.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_iter_raw_docs_missing_index_demotes_readiness(
        self, global_config, embed_func, mock_client
    ):
        mock_client.search = AsyncMock(side_effect=_missing_index_error())

        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()

            batches = [batch async for batch in s._iter_raw_docs(batch_size=2)]

            assert batches == []
            assert s._index_ready is False
            mock_client.create_pit.assert_awaited_once()
            mock_client.delete_pit.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_finalize(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch.object(
                ClientManager, "release_client", new_callable=AsyncMock
            ) as mock_release:
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.finalize()
                mock_release.assert_awaited_once()
                assert s.client is None


# ---------------------------------------------------------------------------
# DocStatus Storage
# ---------------------------------------------------------------------------


class TestDocStatusStorage:
    """Tests for OpenSearchDocStatusStorage including aggregations, pagination, and data normalization."""

    def _make(self, global_config, embed_func, workspace="test"):
        return OpenSearchDocStatusStorage(
            namespace="doc_status",
            global_config=global_config,
            embedding_func=embed_func,
            workspace=workspace,
        )

    @pytest.mark.asyncio
    async def test_index_name(self, global_config, embed_func):
        s = self._make(global_config, embed_func)
        assert s._index_name == "test_doc_status"

    @pytest.mark.asyncio
    async def test_initialize_creates_index(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            mock_client.indices.create.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_get_by_id(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {
                        "_id": "doc-abc",
                        "found": True,
                        "_source": {"status": "processed", "file_path": "/a.txt"},
                    }
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            doc = await s.get_by_id("doc-abc")
            assert doc["status"] == "processed"
            assert doc["_id"] == "doc-abc"
            mock_client.mget.assert_awaited_once_with(
                index=s._index_name, body={"ids": ["doc-abc"]}
            )

    @pytest.mark.asyncio
    async def test_get_by_id_not_found(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={"docs": [{"_id": "missing", "found": False}]}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.get_by_id("missing") is None
            mock_client.get.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_upsert_sets_chunks_list_default(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (1, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.upsert({"d1": {"status": "pending"}})
                actions = mock_bulk.call_args[0][1]
                assert actions[0]["_source"]["chunks_list"] == []

    @pytest.mark.asyncio
    async def test_get_status_counts(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {"hits": [], "total": {"value": 0}},
                "aggregations": {
                    "status_counts": {
                        "buckets": [
                            {"key": "processed", "doc_count": 3},
                            {"key": "pending", "doc_count": 1},
                        ]
                    }
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            counts = await s.get_status_counts()
            assert counts == {"processed": 3, "pending": 1}

    @pytest.mark.asyncio
    async def test_get_all_status_counts_includes_all(
        self, global_config, embed_func, mock_client
    ):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {"hits": [], "total": {"value": 0}},
                "aggregations": {
                    "status_counts": {
                        "buckets": [
                            {"key": "processed", "doc_count": 5},
                            {"key": "failed", "doc_count": 2},
                        ]
                    }
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            counts = await s.get_all_status_counts()
            assert counts["all"] == 7
            assert counts["processed"] == 5

    @pytest.mark.asyncio
    async def test_get_docs_by_status(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "d1",
                            "_source": {
                                "status": "processed",
                                "file_path": "/a.txt",
                                "content_summary": "s",
                                "content_length": 10,
                                "chunks_count": 1,
                                "created_at": 100,
                                "updated_at": 200,
                            },
                            "sort": ["d1"],
                        },
                    ],
                    "total": {"value": 1},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.get_docs_by_status(DocStatus.PROCESSED)
            assert "d1" in result
            assert isinstance(result["d1"], DocProcessingStatus)

    @pytest.mark.asyncio
    async def test_get_docs_paginated(self, global_config, embed_func, mock_client):
        """Page 1 returns results directly without search_after."""
        mock_client.count = AsyncMock(return_value={"count": 50})
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "d1",
                            "_source": {
                                "status": "processed",
                                "file_path": "/a.txt",
                                "content_summary": "s",
                                "content_length": 10,
                                "chunks_count": 1,
                                "created_at": 100,
                                "updated_at": 200,
                            },
                            "sort": [200, "d1"],
                        },
                    ],
                    "total": {"value": 50},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            docs, total = await s.get_docs_paginated(page=1, page_size=10)
            assert total == 50
            assert len(docs) == 1
            assert docs[0][0] == "d1"
            # Page 1: no search_after needed, single search call
            assert mock_client.search.await_count == 1
            body = mock_client.search.call_args.kwargs.get(
                "body"
            ) or mock_client.search.call_args[1].get("body", {})
            assert "search_after" not in body

    @pytest.mark.asyncio
    async def test_get_docs_paginated_page2_uses_search_after(
        self, global_config, embed_func, mock_client
    ):
        """Page 2 skips page 1 results via search_after."""
        mock_client.count = AsyncMock(return_value={"count": 50})
        call_count = {"n": 0}

        async def search_side_effect(*args, **kwargs):
            call_count["n"] += 1
            body = kwargs.get("body", {})
            if "search_after" not in body:
                # First call: skip batch
                return {
                    "hits": {
                        "hits": [
                            {
                                "_id": f"skip{i}",
                                "_source": {
                                    "status": "processed",
                                    "file_path": f"/{i}.txt",
                                    "content_summary": "s",
                                    "content_length": 1,
                                    "chunks_count": 1,
                                    "created_at": 100,
                                    "updated_at": 100 + i,
                                },
                                "sort": [100 + i, f"skip{i}"],
                            }
                            for i in range(10)
                        ],
                        "total": {"value": 50},
                    }
                }
            else:
                # Second call: actual page
                return {
                    "hits": {
                        "hits": [
                            {
                                "_id": "page2_doc",
                                "_source": {
                                    "status": "pending",
                                    "file_path": "/p2.txt",
                                    "content_summary": "s",
                                    "content_length": 1,
                                    "chunks_count": 1,
                                    "created_at": 200,
                                    "updated_at": 300,
                                },
                                "sort": [300, "page2_doc"],
                            }
                        ],
                        "total": {"value": 50},
                    }
                }

        mock_client.search = AsyncMock(side_effect=search_side_effect)
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            docs, total = await s.get_docs_paginated(page=2, page_size=10)
            assert total == 50
            assert len(docs) == 1
            assert docs[0][0] == "page2_doc"
            # 2 search calls: 1 skip + 1 fetch
            assert mock_client.search.await_count == 2

    @pytest.mark.asyncio
    async def test_get_docs_paginated_empty_index(
        self, global_config, embed_func, mock_client
    ):
        """Empty index returns empty list with total 0."""
        mock_client.count = AsyncMock(return_value={"count": 0})
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            docs, total = await s.get_docs_paginated(page=1, page_size=10)
            assert total == 0
            assert docs == []
            mock_client.search.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_get_docs_paginated_page_beyond_total(
        self, global_config, embed_func, mock_client
    ):
        """Requesting a page beyond total docs returns empty list."""
        mock_client.count = AsyncMock(return_value={"count": 5})
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            docs, total = await s.get_docs_paginated(page=100, page_size=10)
            assert total == 5
            assert docs == []

    @pytest.mark.asyncio
    async def test_get_docs_paginated_with_status_filter(
        self, global_config, embed_func, mock_client
    ):
        """Status filter is passed as term query."""
        mock_client.count = AsyncMock(return_value={"count": 3})
        mock_client.search = AsyncMock(
            return_value={
                "hits": {"hits": [], "total": {"value": 3}},
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            docs, total = await s.get_docs_paginated(
                status_filter=DocStatus.PROCESSED, page=1, page_size=10
            )
            assert total == 3
            # Verify count query used the status filter
            count_body = mock_client.count.call_args.kwargs.get("body", {})
            assert count_body["query"] == {"term": {"status": "processed"}}

    @pytest.mark.asyncio
    async def test_get_doc_by_file_path(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "d1",
                            "_source": {
                                "file_path": "/test.txt",
                                "status": "processed",
                            },
                        },
                    ],
                    "total": {"value": 1},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            doc = await s.get_doc_by_file_path("/test.txt")
            assert doc is not None
            assert doc["_id"] == "d1"

    @pytest.mark.asyncio
    async def test_get_doc_by_file_path_not_found(
        self, global_config, embed_func, mock_client
    ):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {"hits": [], "total": {"value": 0}},
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.get_doc_by_file_path("/nope.txt") is None

    @pytest.mark.asyncio
    async def test_prepare_doc_status_data(self, global_config, embed_func):
        s = self._make(global_config, embed_func)
        raw = {"_id": "x", "status": "processed", "error": "oops"}
        data = s._prepare_doc_status_data(raw)
        assert "_id" not in data
        assert data["error_msg"] == "oops"
        assert "error" not in data
        assert data["file_path"] == "no-file-path"
        assert data["metadata"] == {}

    @pytest.mark.asyncio
    async def test_drop_error_marks_index_not_ready_and_next_upsert_recreates_index(
        self, global_config, embed_func, mock_client
    ):
        mock_client.indices.delete = AsyncMock(
            side_effect=OpenSearchException("drop failed")
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (1, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                with patch.object(
                    s, "_create_index_if_not_exists", new_callable=AsyncMock
                ) as mock_create:
                    result = await s.drop()
                    assert result["status"] == "error"
                    assert s._index_ready is False
                    await s.upsert({"d1": {"status": "pending"}})
                    mock_create.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_upsert_after_drop_recreates_index(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (1, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                with patch.object(
                    s, "_create_index_if_not_exists", new_callable=AsyncMock
                ) as mock_create:
                    await s.drop()
                    await s.upsert({"d1": {"status": "pending"}})
                    mock_create.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_reads_short_circuit_after_drop(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.drop()

            assert await s.get_all_status_counts() == {}
            assert await s.get_docs_paginated(page=1, page_size=10) == ([], 0)
            assert await s.get_doc_by_file_path("/a.txt") is None
            assert await s.get_docs_by_status(DocStatus.PROCESSED) == {}

            mock_client.count.assert_not_awaited()
            mock_client.search.assert_not_awaited()
            mock_client.create_pit.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_read_missing_index_demotes_readiness(
        self, global_config, embed_func, mock_client
    ):
        mock_client.search = AsyncMock(side_effect=_missing_index_error())
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()

            assert await s.get_all_status_counts() == {}
            assert await s.get_all_status_counts() == {}
            assert s._index_ready is False
            assert mock_client.search.await_count == 1


# ---------------------------------------------------------------------------
# Graph Storage
# ---------------------------------------------------------------------------


class TestGraphStorage:
    """Tests for OpenSearchGraphStorage node/edge CRUD, batch ops, BFS, and label queries."""

    def _make(self, global_config, embed_func, workspace="test"):
        return OpenSearchGraphStorage(
            namespace="chunk_entity_relation",
            global_config=global_config,
            embedding_func=embed_func,
            workspace=workspace,
        )

    @pytest.mark.asyncio
    async def test_index_names(self, global_config, embed_func):
        s = self._make(global_config, embed_func)
        assert s._nodes_index == "test_chunk_entity_relation-nodes"
        assert s._edges_index == "test_chunk_entity_relation-edges"

    @pytest.mark.asyncio
    async def test_initialize_creates_both_indices(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert mock_client.indices.create.await_count == 2

    @pytest.mark.asyncio
    async def test_has_node_true(self, global_config, embed_func, mock_client):
        mock_client.exists = AsyncMock(return_value=True)
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.has_node("Alice") is True

    @pytest.mark.asyncio
    async def test_has_node_false(self, global_config, embed_func, mock_client):
        mock_client.exists = AsyncMock(return_value=False)
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.has_node("Nobody") is False

    @pytest.mark.asyncio
    async def test_has_edge(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {"hits": [], "total": {"value": 1}},
                "aggregations": {
                    "status_counts": {"buckets": []},
                    "src": {"buckets": []},
                    "tgt": {"buckets": []},
                    "source_degrees": {"buckets": []},
                    "target_degrees": {"buckets": []},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.has_edge("A", "B") is True

    @pytest.mark.asyncio
    async def test_node_degree(self, global_config, embed_func, mock_client):
        mock_client.count = AsyncMock(return_value={"count": 3})
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.node_degree("A") == 3

    @pytest.mark.asyncio
    async def test_get_node(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {
                        "_id": "Alice",
                        "found": True,
                        "_source": {
                            "entity_type": "person",
                            "description": "A researcher",
                        },
                    }
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            node = await s.get_node("Alice")
            assert node["entity_type"] == "person"
            assert node["_id"] == "Alice"
            mock_client.mget.assert_awaited_once_with(
                index=s._nodes_index, body={"ids": ["Alice"]}
            )

    @pytest.mark.asyncio
    async def test_get_node_not_found(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={"docs": [{"_id": "Nobody", "found": False}]}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.get_node("Nobody") is None
            mock_client.get.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_get_edge(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "e1",
                            "_source": {
                                "source_node_id": "A",
                                "target_node_id": "B",
                                "weight": 1.0,
                            },
                        },
                    ],
                    "total": {"value": 1},
                },
                "aggregations": {
                    "status_counts": {"buckets": []},
                    "src": {"buckets": []},
                    "tgt": {"buckets": []},
                    "source_degrees": {"buckets": []},
                    "target_degrees": {"buckets": []},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            edge = await s.get_edge("A", "B")
            assert edge is not None
            assert edge["weight"] == 1.0

    @pytest.mark.asyncio
    async def test_get_node_edges(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "e1",
                            "_source": {"source_node_id": "A", "target_node_id": "B"},
                            "sort": [1],
                        },
                        {
                            "_id": "e2",
                            "_source": {"source_node_id": "C", "target_node_id": "A"},
                            "sort": [2],
                        },
                    ],
                    "total": {"value": 2},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            edges = await s.get_node_edges("A")
            assert len(edges) == 2
            assert ("A", "B") in edges

    @pytest.mark.asyncio
    async def test_get_nodes_batch(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {"_id": "A", "found": True, "_source": {"entity_type": "person"}},
                    {"_id": "B", "found": False},
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.get_nodes_batch(["A", "B"])
            assert "A" in result
            assert "B" not in result

    @pytest.mark.asyncio
    async def test_node_degrees_batch(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {"hits": [], "total": {"value": 0}},
                "aggregations": {
                    "source_degrees": {"buckets": [{"key": "A", "doc_count": 2}]},
                    "target_degrees": {
                        "buckets": [
                            {"key": "A", "doc_count": 1},
                            {"key": "B", "doc_count": 3},
                        ]
                    },
                    "status_counts": {"buckets": []},
                    "src": {"buckets": []},
                    "tgt": {"buckets": []},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            degrees = await s.node_degrees_batch(["A", "B"])
            assert degrees["A"] == 3  # 2 + 1
            assert degrees["B"] == 3

    @pytest.mark.asyncio
    async def test_upsert_node(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.upsert_node(
                "Alice", {"entity_type": "person", "source_id": "c1<SEP>c2"}
            )
            mock_client.index.assert_awaited()
            call_kwargs = mock_client.index.call_args
            assert call_kwargs.kwargs["id"] == "Alice"
            body = call_kwargs.kwargs["body"]
            assert body["source_ids"] == ["c1", "c2"]
            assert body["entity_id"] == "Alice"

    @pytest.mark.asyncio
    async def test_upsert_edge(self, global_config, embed_func, mock_client):
        mock_client.exists = AsyncMock(return_value=False)
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.upsert_edge("A", "B", {"weight": "1.0", "description": "knows"})
            # Should call index twice: once for ensuring source node, once for edge
            assert mock_client.index.await_count == 2

    @pytest.mark.asyncio
    async def test_upsert_after_drop_recreates_indices(
        self, global_config, embed_func, mock_client
    ):
        mock_client.exists = AsyncMock(return_value=False)
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            with patch.object(
                s, "_create_indices_if_not_exist", new_callable=AsyncMock
            ) as mock_create:
                await s.initialize()
                mock_create.reset_mock()
                await s.drop()
                await s.upsert_edge("A", "B", {"weight": "1.0"})
                mock_create.assert_awaited_once()
                assert mock_client.index.await_count == 2

    @pytest.mark.asyncio
    async def test_reads_short_circuit_after_drop(
        self, global_config, embed_func, mock_client
    ):
        mock_client.transport = AsyncMock()
        mock_client.transport.perform_request = AsyncMock(
            side_effect=Exception("PPL not available")
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.drop()

            graph = await s.get_knowledge_graph("A", max_depth=2)

            assert await s.get_node("A") is None
            assert await s.get_all_labels() == []
            assert await s.has_edge("A", "B") is False
            assert await s.node_degree("A") == 0
            assert len(graph.nodes) == 0
            assert len(graph.edges) == 0

            mock_client.mget.assert_not_awaited()
            mock_client.search.assert_not_awaited()
            mock_client.create_pit.assert_not_awaited()
            mock_client.count.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_read_missing_index_demotes_readiness(
        self, global_config, embed_func, mock_client
    ):
        mock_client.transport = AsyncMock()
        mock_client.transport.perform_request = AsyncMock(
            side_effect=Exception("PPL not available")
        )
        mock_client.mget = AsyncMock(side_effect=_missing_index_error())
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()

            assert await s.get_node("A") is None
            assert await s.get_node("A") is None
            assert s._indices_ready is False
            assert mock_client.mget.await_count == 1

    @pytest.mark.asyncio
    async def test_delete_node(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.delete_node("Alice")
            mock_client.delete_by_query.assert_awaited_once()
            mock_client.delete.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_remove_nodes(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (2, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.remove_nodes(["A", "B"])
                mock_client.delete_by_query.assert_awaited_once()
                mock_bulk.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_remove_edges(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.remove_edges([("A", "B"), ("C", "D")])
            mock_client.delete_by_query.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_get_all_labels(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {"_id": "Alice", "sort": ["Alice"]},
                        {"_id": "Bob", "sort": ["Bob"]},
                    ],
                    "total": {"value": 2},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            labels = await s.get_all_labels()
            assert labels == ["Alice", "Bob"]

    @pytest.mark.asyncio
    async def test_get_popular_labels(self, global_config, embed_func, mock_client):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {"hits": [], "total": {"value": 0}},
                "aggregations": {
                    "src": {
                        "buckets": [
                            {"key": "A", "doc_count": 5},
                            {"key": "B", "doc_count": 2},
                        ]
                    },
                    "tgt": {"buckets": [{"key": "A", "doc_count": 3}]},
                    "status_counts": {"buckets": []},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            labels = await s.get_popular_labels(limit=10)
            assert labels[0] == "A"  # degree 8 > B degree 2

    @pytest.mark.asyncio
    async def test_search_labels_empty_query(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.search_labels("") == []

    @pytest.mark.asyncio
    async def test_drop(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.drop()
            assert result["status"] == "success"
            assert mock_client.indices.delete.await_count == 2

    @pytest.mark.asyncio
    async def test_drop_partial_error_marks_indices_not_ready_and_next_upsert_recreates_indices(
        self, global_config, embed_func, mock_client
    ):
        mock_client.exists = AsyncMock(return_value=False)
        mock_client.indices.delete = AsyncMock(
            side_effect=[None, OpenSearchException("edges drop failed")]
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            with patch.object(
                s, "_create_indices_if_not_exist", new_callable=AsyncMock
            ) as mock_create:
                result = await s.drop()
                assert result["status"] == "error"
                assert "edges drop failed" in result["message"]
                assert s._indices_ready is False
                await s.upsert_edge("A", "B", {"weight": "1.0"})
                mock_create.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_drop_treats_missing_graph_indices_as_success(
        self, global_config, embed_func, mock_client
    ):
        mock_client.indices.delete = AsyncMock(
            side_effect=[_missing_index_error(), _missing_index_error()]
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.drop()
            assert result["status"] == "success"
            assert s._indices_ready is False

    @pytest.mark.asyncio
    async def test_construct_graph_node(self, global_config, embed_func):
        s = self._make(global_config, embed_func)
        node = s._construct_graph_node(
            "Alice",
            {
                "entity_type": "person",
                "description": "A researcher",
                "_id": "Alice",
                "entity_id": "Alice",
            },
        )
        assert node.id == "Alice"
        assert "entity_type" in node.properties
        assert "_id" not in node.properties
        assert "entity_id" not in node.properties

    @pytest.mark.asyncio
    async def test_construct_graph_edge(self, global_config, embed_func):
        s = self._make(global_config, embed_func)
        edge = s._construct_graph_edge(
            "e1",
            {
                "source_node_id": "A",
                "target_node_id": "B",
                "relationship": "knows",
                "weight": 1.0,
            },
        )
        assert edge.source == "A"
        assert edge.target == "B"
        assert edge.type == "knows"
        assert "source_node_id" not in edge.properties

    @pytest.mark.asyncio
    async def test_bfs_subgraph_start_not_found(
        self, global_config, embed_func, mock_client
    ):
        mock_client.mget = AsyncMock(
            return_value={"docs": [{"_id": "NonExistent", "found": False}]}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.get_knowledge_graph("NonExistent", max_depth=2)
            assert len(result.nodes) == 0
            assert len(result.edges) == 0


class TestGraphPPLDetection:
    """Tests for PPL graphlookup detection and server-side BFS."""

    def _make(self, global_config, embed_func, workspace="test"):
        return OpenSearchGraphStorage(
            namespace="chunk_entity_relation",
            global_config=global_config,
            embedding_func=embed_func,
            workspace=workspace,
        )

    @pytest.mark.asyncio
    async def test_ppl_detected_when_available(
        self, global_config, embed_func, mock_client
    ):
        """When PPL endpoint responds successfully, graphlookup should be detected."""
        mock_client.transport = AsyncMock()
        mock_client.transport.perform_request = AsyncMock(
            return_value={"datarows": [], "schema": []}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert s._ppl_graphlookup_available is True

    @pytest.mark.asyncio
    async def test_ppl_not_detected_when_endpoint_fails(
        self, global_config, embed_func, mock_client
    ):
        """When PPL endpoint fails, should fall back to client-side BFS."""
        mock_client.transport = AsyncMock()
        mock_client.transport.perform_request = AsyncMock(
            side_effect=Exception("PPL not supported")
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert s._ppl_graphlookup_available is False

    @pytest.mark.asyncio
    async def test_env_override_true(self, global_config, embed_func, mock_client):
        with patch.dict("os.environ", {"OPENSEARCH_USE_PPL_GRAPHLOOKUP": "true"}):
            with patch.object(ClientManager, "get_client", return_value=mock_client):
                s = self._make(global_config, embed_func)
                await s.initialize()
                assert s._ppl_graphlookup_available is True
                # Should NOT have called transport.perform_request for detection
                mock_client.transport.perform_request.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_env_override_false(self, global_config, embed_func, mock_client):
        mock_client.transport = AsyncMock()
        mock_client.transport.perform_request = AsyncMock(
            return_value={"datarows": [], "schema": []}
        )
        with patch.dict("os.environ", {"OPENSEARCH_USE_PPL_GRAPHLOOKUP": "false"}):
            with patch.object(ClientManager, "get_client", return_value=mock_client):
                s = self._make(global_config, embed_func)
                await s.initialize()
                assert s._ppl_graphlookup_available is False

    @pytest.mark.asyncio
    async def test_ppl_bfs_calls_ppl_endpoint(
        self, global_config, embed_func, mock_client
    ):
        """When PPL is available, get_knowledge_graph should use PPL endpoint."""
        mock_client.transport = AsyncMock()
        # PPL response: connected_edges contains dicts with source_node_id/target_node_id
        ppl_response = {
            "schema": [
                {"name": "entity_id", "type": "string"},
                {"name": "connected_edges", "type": "struct"},
            ],
            "datarows": [
                [
                    "A",
                    [  # connected_edges array
                        {
                            "source_node_id": "A",
                            "target_node_id": "B",
                            "weight": 1.0,
                            "_depth": 0,
                        },
                        {
                            "source_node_id": "B",
                            "target_node_id": "C",
                            "weight": 0.5,
                            "_depth": 1,
                        },
                    ],
                ]
            ],
        }
        mock_client.transport.perform_request = AsyncMock(return_value=ppl_response)
        # get_node for start node verification
        mock_client.get = AsyncMock(
            return_value={
                "_id": "A",
                "_source": {"entity_type": "person", "description": "Node A"},
            }
        )
        # mget for batch node fetch (only B and C, A is already added)
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {"_id": "B", "found": True, "_source": {"entity_type": "person"}},
                    {"_id": "C", "found": True, "_source": {"entity_type": "person"}},
                ]
            }
        )
        # search for final edge fetch
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "e1",
                            "_source": {
                                "source_node_id": "A",
                                "target_node_id": "B",
                                "relationship": "knows",
                            },
                        },
                        {
                            "_id": "e2",
                            "_source": {
                                "source_node_id": "B",
                                "target_node_id": "C",
                                "relationship": "knows",
                            },
                        },
                    ],
                    "total": {"value": 2},
                },
                "aggregations": {
                    "status_counts": {"buckets": []},
                    "src": {"buckets": []},
                    "tgt": {"buckets": []},
                },
            }
        )

        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert s._ppl_graphlookup_available is True

            result = await s.get_knowledge_graph("A", max_depth=2)
            assert len(result.nodes) == 3
            assert len(result.edges) == 2
            # Verify PPL was called (2 for detection + 1 for actual query)
            assert mock_client.transport.perform_request.await_count == 3
            # Verify the PPL query uses nodes index as source
            actual_query = mock_client.transport.perform_request.call_args_list[2]
            ppl_body = actual_query.kwargs.get("body") or actual_query[1].get(
                "body", {}
            )
            if isinstance(ppl_body, dict):
                assert s._nodes_index in ppl_body.get("query", "")

    @pytest.mark.asyncio
    async def test_ppl_bfs_falls_back_on_query_failure(
        self, global_config, embed_func, mock_client
    ):
        """If PPL query fails at runtime, should fall back to client-side BFS."""
        call_count = {"n": 0}

        async def ppl_side_effect(*args, **kwargs):
            call_count["n"] += 1
            if call_count["n"] <= 2:
                # Detection calls succeed
                return {"datarows": [], "schema": []}
            # Actual query fails
            raise Exception("PPL query timeout")

        mock_client.transport = AsyncMock()
        mock_client.transport.perform_request = AsyncMock(side_effect=ppl_side_effect)
        mock_client.mget = AsyncMock(
            return_value={"docs": [{"_id": "A", "found": False}]}
        )

        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert s._ppl_graphlookup_available is True

            # Should fall back to _bfs_subgraph, which returns empty (node not found)
            result = await s.get_knowledge_graph("A", max_depth=2)
            assert len(result.nodes) == 0

    @pytest.mark.asyncio
    async def test_escape_ppl(self, global_config, embed_func):
        s = self._make(global_config, embed_func)
        assert s._escape_ppl("it's") == "it\\'s"
        assert s._escape_ppl("normal") == "normal"
        assert s._escape_ppl("back\\slash") == "back\\\\slash"
        assert s._escape_ppl("both\\and'quote") == "both\\\\and\\'quote"

    @pytest.mark.asyncio
    async def test_ppl_bfs_depth_zero_returns_start_only(
        self, global_config, embed_func, mock_client
    ):
        """max_depth=0 should return only the start node without PPL query."""
        mock_client.transport = AsyncMock()
        mock_client.transport.perform_request = AsyncMock(
            return_value={"datarows": [], "schema": []}
        )
        mock_client.get = AsyncMock(
            return_value={"_id": "A", "_source": {"entity_type": "person"}}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert s._ppl_graphlookup_available is True
            result = await s.get_knowledge_graph("A", max_depth=0)
            assert len(result.nodes) == 1
            assert result.nodes[0].id == "A"
            assert len(result.edges) == 0
            # PPL query should NOT have been called for the actual traversal (only 2 detection calls)
            assert mock_client.transport.perform_request.await_count == 2

    @pytest.mark.asyncio
    async def test_ppl_bfs_empty_connected_edges(
        self, global_config, embed_func, mock_client
    ):
        """PPL returns no connected edges — should return only start node."""
        mock_client.transport = AsyncMock()
        ppl_response = {
            "schema": [
                {"name": "entity_id", "type": "string"},
                {"name": "connected_edges", "type": "struct"},
            ],
            "datarows": [["A", []]],
        }
        mock_client.transport.perform_request = AsyncMock(return_value=ppl_response)
        mock_client.get = AsyncMock(
            return_value={"_id": "A", "_source": {"entity_type": "person"}}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.get_knowledge_graph("A", max_depth=2)
            assert len(result.nodes) == 1
            assert result.nodes[0].id == "A"

    @pytest.mark.asyncio
    async def test_upsert_node_adds_entity_id(
        self, global_config, embed_func, mock_client
    ):
        """upsert_node should always include entity_id field for PPL compatibility."""
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.upsert_node("TestNode", {"description": "test"})
            body = mock_client.index.call_args.kwargs["body"]
            assert body["entity_id"] == "TestNode"
            assert body["description"] == "test"

    @pytest.mark.asyncio
    async def test_node_degree_uses_count_api(
        self, global_config, embed_func, mock_client
    ):
        """node_degree should use the count API, not search."""
        mock_client.count = AsyncMock(return_value={"count": 7})
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            degree = await s.node_degree("X")
            assert degree == 7
            # Verify count was called on the edges index
            mock_client.count.assert_awaited()
            call_kwargs = mock_client.count.call_args
            assert s._edges_index in str(call_kwargs)


# ---------------------------------------------------------------------------
# Vector Storage
# ---------------------------------------------------------------------------


class TestVectorStorage:
    """Tests for OpenSearchVectorDBStorage k-NN index, embeddings, cosine conversion, and entity deletion."""

    def _make(self, global_config, embed_func, workspace="test"):
        return OpenSearchVectorDBStorage(
            namespace="entities",
            global_config=global_config,
            embedding_func=embed_func,
            workspace=workspace,
            meta_fields={"content", "entity_name", "src_id", "tgt_id"},
        )

    @pytest.mark.asyncio
    async def test_index_name(self, global_config, embed_func):
        s = self._make(global_config, embed_func)
        assert s._index_name == "test_entities"

    @pytest.mark.asyncio
    async def test_cosine_threshold_required(self, embed_func):
        with pytest.raises(ValueError, match="cosine_better_than_threshold"):
            OpenSearchVectorDBStorage(
                namespace="v",
                global_config={
                    "embedding_batch_num": 10,
                    "vector_db_storage_cls_kwargs": {},
                },
                embedding_func=embed_func,
            )

    @pytest.mark.asyncio
    async def test_initialize_creates_knn_index(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            mock_client.indices.create.assert_awaited_once()
            body = mock_client.indices.create.call_args.kwargs["body"]
            assert body["settings"]["index"]["knn"] is True
            assert body["mappings"]["properties"]["vector"]["dimension"] == 128
            assert (
                body["mappings"]["properties"]["vector"]["method"]["engine"] == "lucene"
            )

    @pytest.mark.asyncio
    async def test_upsert_generates_embeddings(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (2, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.upsert(
                    {
                        "v1": {"content": "hello"},
                        "v2": {"content": "world"},
                    }
                )
                actions = mock_bulk.call_args[0][1]
                assert len(actions) == 2
                assert "vector" in actions[0]["_source"]
                assert len(actions[0]["_source"]["vector"]) == 128

    @pytest.mark.asyncio
    async def test_query_cosine_score_conversion(
        self, global_config, embed_func, mock_client
    ):
        """Test that scores are used directly and threshold filtering works."""
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "v1",
                            "_score": 0.85,
                            "_source": {"content": "match", "entity_name": "E1"},
                        },
                    ],
                    "total": {"value": 1},
                },
                "aggregations": {
                    "status_counts": {"buckets": []},
                    "src": {"buckets": []},
                    "tgt": {"buckets": []},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            results = await s.query("test", top_k=5)
            assert len(results) == 1
            assert results[0]["distance"] == 0.85

    @pytest.mark.asyncio
    async def test_query_filters_below_threshold(
        self, global_config, embed_func, mock_client
    ):
        """Low scores should be filtered out."""
        # score 0.15 < threshold 0.2
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {
                            "_id": "v1",
                            "_score": 0.15,
                            "_source": {"content": "weak match"},
                        },
                    ],
                    "total": {"value": 1},
                },
                "aggregations": {
                    "status_counts": {"buckets": []},
                    "src": {"buckets": []},
                    "tgt": {"buckets": []},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            results = await s.query("test", top_k=5)
            assert len(results) == 0

    @pytest.mark.asyncio
    async def test_query_with_provided_embedding(
        self, global_config, embed_func, mock_client
    ):
        mock_client.search = AsyncMock(
            return_value={
                "hits": {
                    "hits": [
                        {"_id": "v1", "_score": 1.0, "_source": {"content": "exact"}},
                    ],
                    "total": {"value": 1},
                },
                "aggregations": {
                    "status_counts": {"buckets": []},
                    "src": {"buckets": []},
                    "tgt": {"buckets": []},
                },
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            vec = np.random.rand(128).astype(np.float32)
            results = await s.query("test", top_k=5, query_embedding=vec)
            assert len(results) == 1
            assert results[0]["distance"] == 1.0

    @pytest.mark.asyncio
    async def test_get_by_id(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {
                        "_id": "v1",
                        "found": True,
                        "_source": {"content": "hello", "vector": [0.1] * 128},
                    }
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            doc = await s.get_by_id("v1")
            assert doc["id"] == "v1"
            assert doc["content"] == "hello"
            mock_client.mget.assert_awaited_once_with(
                index=s._index_name, body={"ids": ["v1"]}
            )

    @pytest.mark.asyncio
    async def test_get_by_id_not_found(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={"docs": [{"_id": "missing", "found": False}]}
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            assert await s.get_by_id("missing") is None
            mock_client.get.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_get_by_ids(self, global_config, embed_func, mock_client):
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {"_id": "v1", "found": True, "_source": {"content": "a"}},
                    {"_id": "v2", "found": True, "_source": {"content": "b"}},
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            docs = await s.get_by_ids(["v1", "v2"])
            assert docs[0]["id"] == "v1"
            assert docs[1]["id"] == "v2"

    @pytest.mark.asyncio
    async def test_get_vectors_by_ids(self, global_config, embed_func, mock_client):
        vec = [0.1] * 128
        mock_client.mget = AsyncMock(
            return_value={
                "docs": [
                    {"_id": "v1", "found": True, "_source": {"vector": vec}},
                    {"_id": "v2", "found": False},
                ]
            }
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.get_vectors_by_ids(["v1", "v2"])
            assert "v1" in result
            assert "v2" not in result
            assert result["v1"] == vec

    @pytest.mark.asyncio
    async def test_delete(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            with patch(
                "lightrag.kg.opensearch_impl.helpers.async_bulk", new_callable=AsyncMock
            ) as mock_bulk:
                mock_bulk.return_value = (2, [])
                s = self._make(global_config, embed_func)
                await s.initialize()
                await s.delete(["v1", "v2"])
                actions = mock_bulk.call_args[0][1]
                assert len(actions) == 2
                assert all(a["_op_type"] == "delete" for a in actions)

    @pytest.mark.asyncio
    async def test_delete_entity(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.delete_entity("Alice")
            mock_client.delete.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_delete_entity_relation(self, global_config, embed_func, mock_client):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            await s.delete_entity_relation("Alice")
            mock_client.delete_by_query.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_drop_recreates_index(self, global_config, embed_func, mock_client):
        # After drop, _create_knn_index_if_not_exists is called again.
        # First call (init): exists=False -> create. Second call (after drop): exists=False -> create again.
        mock_client.indices.exists = AsyncMock(return_value=False)
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.drop()
            assert result["status"] == "success"
            mock_client.indices.delete.assert_awaited_once()
            # create called twice: once during init, once during drop recreate
            assert mock_client.indices.create.await_count == 2

    @pytest.mark.asyncio
    async def test_drop_delete_error_marks_index_not_ready(
        self, global_config, embed_func, mock_client
    ):
        mock_client.indices.delete = AsyncMock(
            side_effect=OpenSearchException("delete failed")
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.drop()
            assert result["status"] == "error"
            assert s._index_ready is False

    @pytest.mark.asyncio
    async def test_drop_recreate_error_marks_index_not_ready(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            with patch.object(
                s,
                "_create_knn_index_if_not_exists",
                new=AsyncMock(side_effect=OpenSearchException("recreate failed")),
            ):
                result = await s.drop()
                assert result["status"] == "error"
                assert s._index_ready is False

    @pytest.mark.asyncio
    async def test_drop_recreates_index_when_missing(
        self, global_config, embed_func, mock_client
    ):
        mock_client.indices.exists = AsyncMock(return_value=False)
        mock_client.indices.delete = AsyncMock(
            side_effect=NotFoundError(404, "not found")
        )
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            result = await s.drop()
            assert result["status"] == "success"
            assert mock_client.indices.create.await_count == 2

    @pytest.mark.asyncio
    async def test_reads_short_circuit_when_index_not_ready(
        self, global_config, embed_func, mock_client
    ):
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()
            s._index_ready = False

            assert await s.query("test", top_k=5) == []
            assert await s.get_by_id("v1") is None
            assert await s.get_vectors_by_ids(["v1"]) == {}

            mock_client.search.assert_not_awaited()
            mock_client.mget.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_read_missing_index_demotes_readiness(
        self, global_config, embed_func, mock_client
    ):
        mock_client.search = AsyncMock(side_effect=_missing_index_error())
        with patch.object(ClientManager, "get_client", return_value=mock_client):
            s = self._make(global_config, embed_func)
            await s.initialize()

            assert await s.query("test", top_k=5) == []
            assert await s.query("test", top_k=5) == []
            assert s._index_ready is False
            assert mock_client.search.await_count == 1


# ---------------------------------------------------------------------------
# Cosine score edge cases
# ---------------------------------------------------------------------------


class TestScoreThreshold:
    """Verify that raw OpenSearch scores are compared directly against threshold."""

    def test_above_threshold(self):
        assert 0.85 >= 0.2

    def test_below_threshold(self):
        assert 0.15 < 0.2

    def test_exact_threshold(self):
        assert 0.2 >= 0.2