library-desk/tests/test_tenant_scoping.py

"""
Offline tenant-isolation unit tests for every HybridRAG leg (mocked clients).

For each retrieval leg (vector, graph, volatile, documents) plus the
enrichment and persistence phases, assert that the tenant-scoped
collection / label / path is used and that cross-tenant access is refused.

Live probes showed /query/hybrid as user=llm_tester returning jpmschweitzer
pages: the root cause was unscoped ingestion (any tenant could ingest any
page id / any path prefix into its own collection) and an unscoped
Document match in search persistence. These tests pin the fixes.
"""

from unittest.mock import AsyncMock, MagicMock

import pytest

from src.models.hybrid_rag import HybridRAGConfig
from src.services.graph_service import GraphService
from src.services.hybrid_rag_service import HybridRAGService
from src.services.ingestion_service import IngestionService
from src.services.vector_service import VectorService
from src.services.volatile_service import VolatileCacheService

TENANT = "llm_tester"
TENANT_COLLECTION = "library_desk_llm_tester"
TENANT_BASE_LABEL = "User_Llm_Tester"
TENANT_DOC_LABEL = "User_Llm_Tester_Document"


# =============================================================================
# Fixtures
# =============================================================================


@pytest.fixture
def mock_qdrant():
    qdrant = MagicMock()
    qdrant.collection_exists = AsyncMock(return_value=True)
    qdrant.ensure_collection = AsyncMock()
    qdrant.search_vectors = AsyncMock(return_value=[])
    qdrant.search_with_expiry_filter = AsyncMock(return_value=[])
    qdrant.upsert_vector = AsyncMock(return_value=True)
    qdrant.upsert_points = AsyncMock(side_effect=lambda collection_name, points: len(points))
    qdrant.scroll_all_points = AsyncMock(return_value=[])
    qdrant.delete_by_ids = AsyncMock(return_value=0)
    qdrant.delete_by_filter = AsyncMock(return_value=0)
    return qdrant


@pytest.fixture
def mock_ollama():
    ollama = MagicMock()
    ollama.embed = AsyncMock(return_value=[0.1] * 768)
    ollama.embed_batch = AsyncMock(
        side_effect=lambda texts, **kw: [[0.1] * 768 for _ in texts]
    )
    ollama.embed_text = AsyncMock(return_value=[0.1] * 768)
    ollama.generate_text = AsyncMock(return_value="{}")
    return ollama


@pytest.fixture
def mock_neo4j():
    neo4j = MagicMock()
    neo4j.execute_query = AsyncMock(return_value=[])
    neo4j.execute_read = AsyncMock(return_value=[])
    neo4j.execute_write = AsyncMock(return_value=[])
    return neo4j


@pytest.fixture
def mock_wiki():
    wiki = MagicMock()
    wiki.get_page = AsyncMock(return_value=None)
    wiki.list_all_pages = AsyncMock(return_value=[])
    wiki.list_pages = AsyncMock(return_value=[])
    return wiki


@pytest.fixture
def vector_service(mock_qdrant, mock_wiki, mock_ollama):
    return VectorService(mock_qdrant, mock_wiki, mock_ollama)


@pytest.fixture
def graph_service(mock_neo4j, mock_wiki):
    return GraphService(mock_neo4j, mock_wiki)


@pytest.fixture
def volatile_service(mock_qdrant, mock_ollama):
    return VolatileCacheService(mock_qdrant, mock_ollama, MagicMock())


@pytest.fixture
def hybrid_service(vector_service, graph_service, volatile_service, mock_ollama):
    settings = MagicMock()
    settings.ollama_llm_model = "test-model"
    settings.vector_similarity_threshold = 0.7
    return HybridRAGService(
        vector_service=vector_service,
        graph_service=graph_service,
        searxng_client=MagicMock(),
        ollama_client=mock_ollama,
        content_extractor=MagicMock(),
        settings=settings,
        volatile_service=volatile_service,
    )


def _all_cypher(mock_neo4j) -> str:
    """Concatenate all Cypher sent to the mocked Neo4j client (reads + writes)."""
    return "\n".join(
        str(call.args[0])
        for mock in (mock_neo4j.execute_query, mock_neo4j.execute_write)
        for call in mock.await_args_list
    )


# =============================================================================
# Vector leg
# =============================================================================


@pytest.mark.unit
class TestVectorLegScoping:
    async def test_search_uses_tenant_collection(self, vector_service, mock_qdrant):
        await vector_service.search(query="q", user=TENANT)

        mock_qdrant.collection_exists.assert_awaited_with(TENANT_COLLECTION)
        assert (
            mock_qdrant.search_vectors.await_args.kwargs["collection_name"]
            == TENANT_COLLECTION
        )

    async def test_delete_page_chunks_uses_tenant_collection(
        self, vector_service, mock_qdrant
    ):
        await vector_service.delete_page_chunks(page_id=1, user=TENANT)

        assert (
            mock_qdrant.delete_by_filter.await_args.kwargs["collection_name"]
            == TENANT_COLLECTION
        )

    async def test_update_from_page_rejects_foreign_namespace(
        self, vector_service, mock_qdrant, mock_wiki
    ):
        """Ingesting another tenant's page into our collection must fail."""
        mock_wiki.get_page = AsyncMock(return_value={
            "id": 42, "title": "T", "path": "users/jpmschweitzer/secret",
            "content": "secret content", "tags": [],
        })

        summary = await vector_service.update_from_page(page_id=42, user=TENANT)

        assert summary.success is False
        assert "outside user" in (summary.error_message or "")
        mock_qdrant.upsert_points.assert_not_awaited()
        mock_qdrant.delete_by_ids.assert_not_awaited()

    async def test_update_from_page_rejects_sibling_prefix_namespace(
        self, vector_service, mock_qdrant, mock_wiki
    ):
        """users/llm_tester2 is NOT inside llm_tester's namespace."""
        mock_wiki.get_page = AsyncMock(return_value={
            "id": 43, "title": "T", "path": "users/llm_tester2/page",
            "content": "content", "tags": [],
        })

        summary = await vector_service.update_from_page(page_id=43, user=TENANT)

        assert summary.success is False
        mock_qdrant.upsert_points.assert_not_awaited()

    async def test_update_from_page_accepts_own_namespace(
        self, vector_service, mock_qdrant, mock_wiki
    ):
        mock_wiki.get_page = AsyncMock(return_value={
            "id": 44, "title": "T", "path": "users/llm_tester/page",
            "content": "hello world", "tags": [],
        })

        summary = await vector_service.update_from_page(page_id=44, user=TENANT)

        assert summary.success is True
        assert (
            mock_qdrant.upsert_points.await_args.kwargs["collection_name"]
            == TENANT_COLLECTION
        )


# =============================================================================
# Graph leg
# =============================================================================


@pytest.mark.unit
class TestGraphLegScoping:
    async def test_search_documents_uses_tenant_labels(
        self, graph_service, mock_neo4j
    ):
        await graph_service.search_documents(
            query="docker", user=TENANT, keywords_data={"core_keywords": ["docker"]}
        )

        cypher = _all_cypher(mock_neo4j)
        assert f"(e:{TENANT_BASE_LABEL})" in cypher
        assert f"(d:{TENANT_DOC_LABEL}:Document)" in cypher

    async def test_update_from_page_rejects_foreign_namespace(
        self, graph_service, mock_neo4j, mock_wiki
    ):
        mock_wiki.get_page = AsyncMock(return_value={
            "id": 42, "title": "T", "path": "users/jpmschweitzer/secret",
            "content": "secret", "tags": [],
        })

        summary = await graph_service.update_from_page(page_id=42, user=TENANT)

        assert summary.success is False
        assert "outside user" in (summary.error_message or "")
        mock_neo4j.execute_query.assert_not_awaited()

    async def test_update_from_page_writes_tenant_labels(
        self, graph_service, mock_neo4j, mock_wiki
    ):
        mock_wiki.get_page = AsyncMock(return_value={
            "id": 44, "title": "T", "path": "users/llm_tester/page",
            "content": "Uses Docker daily", "tags": [],
        })

        summary = await graph_service.update_from_page(page_id=44, user=TENANT)

        assert summary.success is True
        cypher = _all_cypher(mock_neo4j)
        assert TENANT_DOC_LABEL in cypher
        # No unscoped Document writes
        assert "MERGE (d:Document {" not in cypher

    async def test_entity_mention_count_scoped(self, graph_service, mock_neo4j):
        await graph_service._get_entity_mention_count("Docker", "Technology", TENANT)

        cypher = _all_cypher(mock_neo4j)
        assert f"(d:{TENANT_DOC_LABEL}:Document)" in cypher
        assert "MATCH (d:Document)" not in cypher

    async def test_orphan_queries_scoped(self, graph_service, mock_neo4j):
        await graph_service.find_orphan_entities(TENANT)
        await graph_service.purge_orphan_entities(TENANT)

        cypher = _all_cypher(mock_neo4j)
        assert f"(d:{TENANT_DOC_LABEL}:Document)" in cypher
        assert "(d:Document)-[:MENTIONS]" not in cypher

    async def test_cleanup_broken_relationships_scoped(
        self, graph_service, mock_neo4j
    ):
        await graph_service.cleanup_broken_relationships(TENANT)

        cypher = _all_cypher(mock_neo4j)
        assert f"{TENANT_BASE_LABEL}_SearchQuery" in cypher
        assert "MATCH (sq:SearchQuery)" not in cypher


# =============================================================================
# Volatile leg
# =============================================================================


@pytest.mark.unit
class TestVolatileLegScoping:
    async def test_search_uses_tenant_collection(self, volatile_service, mock_qdrant):
        await volatile_service.search(user=TENANT, query="weather")

        mock_qdrant.collection_exists.assert_awaited_with("volatile_llm_tester")
        assert (
            mock_qdrant.search_with_expiry_filter.await_args.kwargs["collection_name"]
            == "volatile_llm_tester"
        )

    async def test_store_uses_tenant_collection(self, volatile_service, mock_qdrant):
        await volatile_service.store(
            user=TENANT, namespace="weather", key="rotterdam", data={"t": 1}
        )

        assert (
            mock_qdrant.upsert_vector.await_args.kwargs["collection_name"]
            == "volatile_llm_tester"
        )

    def test_collection_name_is_sanitized(self, volatile_service):
        """Raw user strings cannot alias/escape the collection scheme."""
        assert volatile_service._collection_name("llm-tester") == "volatile_llm_tester"
        assert volatile_service._collection_name("Evil User!") == "volatile_evil_user"


# =============================================================================
# Documents leg (Paperless chunks in the tenant collection)
# =============================================================================


@pytest.mark.unit
class TestDocumentLegScoping:
    async def test_document_search_uses_tenant_collection(
        self, hybrid_service, mock_qdrant
    ):
        config = HybridRAGConfig(
            enable_vector=False, enable_graph=False, enable_web=False,
            enable_volatile=False, enable_documents=True,
        )

        await hybrid_service._retrieve_parallel("q", TENANT, config, {})

        assert (
            mock_qdrant.search_vectors.await_args.kwargs["collection_name"]
            == TENANT_COLLECTION
        )
        assert (
            mock_qdrant.search_vectors.await_args.kwargs["filter_conditions"]
            == {"doc_type": "document"}
        )


# =============================================================================
# Enrichment phase
# =============================================================================


@pytest.mark.unit
class TestEnrichmentScoping:
    async def test_related_documents_scoped_to_tenant(
        self, graph_service, mock_neo4j
    ):
        await graph_service.get_related_documents(page_id=1, user=TENANT)

        cypher = _all_cypher(mock_neo4j)
        assert f"(d1:{TENANT_DOC_LABEL}:Document" in cypher
        assert f"(d2:{TENANT_DOC_LABEL}:Document)" in cypher

    async def test_batched_related_documents_scoped_to_tenant(
        self, graph_service, mock_neo4j
    ):
        await graph_service.get_related_documents_batch(
            page_ids=[1, 2], user=TENANT
        )

        cypher = _all_cypher(mock_neo4j)
        assert "UNWIND $page_ids" in cypher
        assert f"(d1:{TENANT_DOC_LABEL}:Document" in cypher
        assert f"(d2:{TENANT_DOC_LABEL}:Document)" in cypher


# =============================================================================
# Persistence phase (SearchQuery / FOUND / WebResult)
# =============================================================================


@pytest.mark.unit
class TestPersistenceScoping:
    async def test_persist_scopes_searchquery_links_and_webresults(
        self, hybrid_service, mock_neo4j
    ):
        final_results = [
            {
                "result": {"page_id": 42, "title": "wiki hit"},
                "source_type": "wiki", "sources": ["vector"],
                "rrf_score": 0.5, "final_rank": 1,
            },
            {
                "result": {"url": "http://example.com", "title": "web hit",
                           "content": "c"},
                "source_type": "web", "sources": ["web"],
                "rrf_score": 0.4, "final_rank": 2,
            },
        ]

        await hybrid_service._persist_search_for_librarian(
            search_id="sid-1", query="q", user=TENANT, keywords_data={},
            raw_results={}, final_results=final_results, timing={},
        )

        # Persistence is now ONE atomic write transaction
        mock_neo4j.execute_write.assert_awaited_once()
        cypher = str(mock_neo4j.execute_write.await_args.args[0])
        # SearchQuery + WebResult nodes carry the tenant label
        assert f"{TENANT_BASE_LABEL}_SearchQuery" in cypher
        assert f"{TENANT_BASE_LABEL}_WebResult" in cypher
        # FOUND link matches only this tenant's Document nodes
        assert f"MATCH (d:{TENANT_DOC_LABEL}:Document {{page_id: link.page_id}})" in cypher
        # The old cross-tenant match must be gone
        assert "MATCH (d:Document {page_id:" not in cypher


# =============================================================================
# Ingestion (feeds the vector/graph legs)
# =============================================================================


@pytest.mark.unit
class TestIngestAllPagesScoping:
    @pytest.fixture
    def ingestion_service(self, vector_service, graph_service, mock_wiki):
        return IngestionService(vector_service, graph_service, mock_wiki)

    async def test_defaults_to_own_namespace(self, ingestion_service, mock_wiki):
        await ingestion_service.ingest_all_pages(user=TENANT)

        mock_wiki.list_all_pages.assert_awaited_once_with(
            path_prefix="users/llm_tester"
        )

    async def test_rejects_foreign_prefix(self, ingestion_service, mock_wiki):
        with pytest.raises(ValueError, match="outside user"):
            await ingestion_service.ingest_all_pages(
                user=TENANT, path_prefix="users/jpmschweitzer"
            )

        mock_wiki.list_all_pages.assert_not_awaited()

    async def test_rejects_sibling_prefix(self, ingestion_service, mock_wiki):
        with pytest.raises(ValueError, match="outside user"):
            await ingestion_service.ingest_all_pages(
                user=TENANT, path_prefix="users/llm_tester2"
            )

    async def test_accepts_subtree_of_own_namespace(
        self, ingestion_service, mock_wiki
    ):
        await ingestion_service.ingest_all_pages(
            user=TENANT, path_prefix="users/llm_tester/tech"
        )

        mock_wiki.list_all_pages.assert_awaited_once_with(
            path_prefix="users/llm_tester/tech"
        )