Every vector call ran on the sync QdrantClient inside async wrapper methods, blocking the FastAPI event loop per Qdrant round-trip. The wrapper now holds an AsyncQdrantClient (timeout via QDRANT_TIMEOUT, default 30s) and awaits all client calls; the wrapper API is unchanged. Call sites off the wrapper were fixed too: the HybridRAG document leg now uses the async search_vectors wrapper instead of the deprecated raw client.search (also fixing its call to the nonexistent ollama.embed_text which made the leg permanently report 'failed'), the health check awaits get_collections, and document_sync's raw delete/upsert calls are awaited (routed through wrappers in the next commit). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QbFZyDvYksazX6nYQYZ67L
415 lines
15 KiB
Python
415 lines
15 KiB
Python
"""
|
|
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.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_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=[])
|
|
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."""
|
|
return "\n".join(
|
|
str(call.args[0]) for call in mock_neo4j.execute_query.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_vector.assert_not_awaited()
|
|
mock_qdrant.delete_by_filter.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_vector.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_vector.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
|
|
|
|
|
|
# =============================================================================
|
|
# 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(
|
|
query="q", user=TENANT, keywords_data={},
|
|
raw_results={}, final_results=final_results, timing={},
|
|
)
|
|
|
|
cypher = _all_cypher(mock_neo4j)
|
|
# 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: $page_id}})" in cypher
|
|
# The old cross-tenant match must be gone
|
|
assert "MATCH (d:Document {page_id: $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"
|
|
)
|