703 lines
24 KiB
Python
703 lines
24 KiB
Python
"""
|
|
Comprehensive End-to-End Multi-Tenant Testing Suite
|
|
|
|
This module tests the complete multi-tenant architecture from API layer through
|
|
storage backends, ensuring proper data isolation, context propagation, and
|
|
composite key enforcement across all operations.
|
|
|
|
Test Categories:
|
|
1. Tenant & KB Management: Creation, retrieval, deletion
|
|
2. Data Isolation: Cross-tenant prevention, composite key enforcement
|
|
3. Document Operations: Upload, retrieval per tenant/KB scope
|
|
4. Query Operations: Entity/relation queries with tenant isolation
|
|
5. Cache Operations: Redis namespace isolation per tenant
|
|
6. Edge Cases: Boundary conditions, error handling, concurrent access
|
|
"""
|
|
|
|
import pytest
|
|
import asyncio
|
|
import os
|
|
from typing import Dict
|
|
from unittest.mock import AsyncMock
|
|
|
|
|
|
from lightrag.models.tenant import Tenant, TenantContext, KnowledgeBase
|
|
from lightrag.services.tenant_service import TenantService
|
|
from lightrag.base import BaseKVStorage
|
|
from lightrag.kg.postgres_tenant_support import get_composite_key
|
|
from lightrag.kg.redis_tenant_support import RedisTenantHelper
|
|
|
|
|
|
# ============================================================================
|
|
# Test Fixtures & Setup
|
|
# ============================================================================
|
|
|
|
|
|
@pytest.fixture
|
|
def testing_mode():
|
|
"""Return current testing mode"""
|
|
return os.getenv("MULTITENANT_MODE", "demo")
|
|
|
|
|
|
@pytest.fixture
|
|
def is_demo_mode(testing_mode):
|
|
"""Check if in demo mode (2 tenants)"""
|
|
return testing_mode == "demo"
|
|
|
|
|
|
@pytest.fixture
|
|
def is_multi_tenant_mode(testing_mode):
|
|
"""Check if multi-tenancy is enabled"""
|
|
return testing_mode in ["demo", "on"]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_kv_storage():
|
|
"""Create mock KV storage for testing"""
|
|
storage = AsyncMock(spec=BaseKVStorage)
|
|
storage.upsert = AsyncMock()
|
|
storage.get = AsyncMock()
|
|
storage.delete = AsyncMock()
|
|
storage.query = AsyncMock()
|
|
return storage
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_tenant_service(mock_kv_storage):
|
|
"""Create mock tenant service"""
|
|
service = TenantService(mock_kv_storage)
|
|
return service
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_tenants() -> Dict[str, Dict]:
|
|
"""Sample tenants for multi-tenant testing"""
|
|
return {
|
|
"tenant_a": {
|
|
"tenant_id": "tenant-a",
|
|
"name": "Tenant A",
|
|
"description": "Test Tenant A",
|
|
},
|
|
"tenant_b": {
|
|
"tenant_id": "tenant-b",
|
|
"name": "Tenant B",
|
|
"description": "Test Tenant B",
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_kbs() -> Dict[str, Dict]:
|
|
"""Sample knowledge bases"""
|
|
return {
|
|
"kb_a1": {
|
|
"kb_id": "kb-a-1",
|
|
"tenant_id": "tenant-a",
|
|
"name": "KB A-1",
|
|
"description": "Knowledge Base A-1",
|
|
},
|
|
"kb_a2": {
|
|
"kb_id": "kb-a-2",
|
|
"tenant_id": "tenant-a",
|
|
"name": "KB A-2",
|
|
"description": "Knowledge Base A-2",
|
|
},
|
|
"kb_b1": {
|
|
"kb_id": "kb-b-1",
|
|
"tenant_id": "tenant-b",
|
|
"name": "KB B-1",
|
|
"description": "Knowledge Base B-1",
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_documents() -> Dict[str, Dict]:
|
|
"""Sample documents for testing"""
|
|
return {
|
|
"doc_a1_1": {
|
|
"doc_id": "doc-a1-1",
|
|
"tenant_id": "tenant-a",
|
|
"kb_id": "kb-a-1",
|
|
"title": "Document A1-1",
|
|
"content": "Content for tenant A, KB 1",
|
|
"status": "active",
|
|
},
|
|
"doc_a1_2": {
|
|
"doc_id": "doc-a1-2",
|
|
"tenant_id": "tenant-a",
|
|
"kb_id": "kb-a-1",
|
|
"title": "Document A1-2",
|
|
"content": "Another document for tenant A, KB 1",
|
|
"status": "active",
|
|
},
|
|
"doc_a2_1": {
|
|
"doc_id": "doc-a2-1",
|
|
"tenant_id": "tenant-a",
|
|
"kb_id": "kb-a-2",
|
|
"title": "Document A2-1",
|
|
"content": "Document for tenant A, KB 2",
|
|
"status": "active",
|
|
},
|
|
"doc_b1_1": {
|
|
"doc_id": "doc-b1-1",
|
|
"tenant_id": "tenant-b",
|
|
"kb_id": "kb-b-1",
|
|
"title": "Document B1-1",
|
|
"content": "Content for tenant B, KB 1",
|
|
"status": "active",
|
|
},
|
|
}
|
|
|
|
|
|
# ============================================================================
|
|
# Composite Key Pattern Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestCompositeKeyPattern:
|
|
"""Test composite key generation and enforcement"""
|
|
|
|
def test_composite_key_generation(self):
|
|
"""Test basic composite key generation"""
|
|
key = get_composite_key("tenant-a", "kb-1", "doc-123")
|
|
assert key == "tenant-a:kb-1:doc-123"
|
|
assert key.count(":") == 2
|
|
|
|
def test_composite_key_with_special_chars(self):
|
|
"""Test composite key with special characters"""
|
|
key = get_composite_key("tenant_a", "kb-prod_v2", "entity_id")
|
|
assert key == "tenant_a:kb-prod_v2:entity_id"
|
|
|
|
def test_composite_key_uniqueness(self):
|
|
"""Test that different tenant/kb combos create different keys"""
|
|
key1 = get_composite_key("tenant-a", "kb-1", "doc-123")
|
|
key2 = get_composite_key("tenant-b", "kb-1", "doc-123")
|
|
key3 = get_composite_key("tenant-a", "kb-2", "doc-123")
|
|
|
|
assert key1 != key2
|
|
assert key1 != key3
|
|
assert key2 != key3
|
|
|
|
def test_composite_key_deterministic(self):
|
|
"""Test that composite key generation is deterministic"""
|
|
key1 = get_composite_key("tenant-a", "kb-1", "doc-123")
|
|
key2 = get_composite_key("tenant-a", "kb-1", "doc-123")
|
|
|
|
assert key1 == key2
|
|
|
|
|
|
# ============================================================================
|
|
# Data Isolation Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestDataIsolation:
|
|
"""Test multi-tenant data isolation at storage layer"""
|
|
|
|
def test_tenant_a_cannot_access_tenant_b_docs(self, sample_documents):
|
|
"""Test that Tenant A cannot access Tenant B documents"""
|
|
# Simulate storage query
|
|
tenant_a_docs = [
|
|
d for d in sample_documents.values() if d["tenant_id"] == "tenant-a"
|
|
]
|
|
tenant_b_docs = [
|
|
d for d in sample_documents.values() if d["tenant_id"] == "tenant-b"
|
|
]
|
|
|
|
assert len(tenant_a_docs) == 3
|
|
assert len(tenant_b_docs) == 1
|
|
assert all(d["tenant_id"] == "tenant-a" for d in tenant_a_docs)
|
|
assert all(d["tenant_id"] == "tenant-b" for d in tenant_b_docs)
|
|
assert (
|
|
len(
|
|
set(d["doc_id"] for d in tenant_a_docs)
|
|
& set(d["doc_id"] for d in tenant_b_docs)
|
|
)
|
|
== 0
|
|
)
|
|
|
|
def test_kb_isolation_within_same_tenant(self, sample_documents):
|
|
"""Test KB-level isolation within same tenant"""
|
|
# Get docs for tenant-a, kb-a-1
|
|
kb_a1_docs = [
|
|
d
|
|
for d in sample_documents.values()
|
|
if d["tenant_id"] == "tenant-a" and d["kb_id"] == "kb-a-1"
|
|
]
|
|
kb_a2_docs = [
|
|
d
|
|
for d in sample_documents.values()
|
|
if d["tenant_id"] == "tenant-a" and d["kb_id"] == "kb-a-2"
|
|
]
|
|
|
|
assert len(kb_a1_docs) == 2
|
|
assert len(kb_a2_docs) == 1
|
|
assert (
|
|
len(
|
|
set(d["doc_id"] for d in kb_a1_docs)
|
|
& set(d["doc_id"] for d in kb_a2_docs)
|
|
)
|
|
== 0
|
|
)
|
|
|
|
def test_composite_key_prevents_id_collision(self):
|
|
"""Test that composite keys prevent ID collisions across tenants"""
|
|
# Same doc_id in different tenant/KB combos should be different
|
|
key_a1 = get_composite_key("tenant-a", "kb-1", "doc-123")
|
|
key_a2 = get_composite_key("tenant-a", "kb-2", "doc-123")
|
|
key_b1 = get_composite_key("tenant-b", "kb-1", "doc-123")
|
|
|
|
# All keys are unique even though doc-123 is same
|
|
assert len({key_a1, key_a2, key_b1}) == 3
|
|
|
|
|
|
# ============================================================================
|
|
# Redis Namespace Isolation Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestRedisNamespaceIsolation:
|
|
"""Test Redis key prefixing for tenant isolation"""
|
|
|
|
def test_redis_tenant_key_generation(self):
|
|
"""Test Redis tenant-scoped key generation"""
|
|
key = RedisTenantHelper.make_tenant_key("tenant-a", "kb-1", "cache:user:123")
|
|
assert key == "tenant-a:kb-1:cache:user:123"
|
|
|
|
def test_redis_tenant_key_pattern(self):
|
|
"""Test Redis tenant key pattern matching"""
|
|
pattern = RedisTenantHelper.get_tenant_key_pattern("tenant-a", "kb-1")
|
|
assert pattern == "tenant-a:kb-1:*"
|
|
|
|
def test_redis_tenant_key_custom_pattern(self):
|
|
"""Test custom pattern with tenant scope"""
|
|
pattern = RedisTenantHelper.get_tenant_key_pattern(
|
|
"tenant-a", "kb-1", "cache:*"
|
|
)
|
|
assert pattern == "tenant-a:kb-1:cache:*"
|
|
|
|
def test_redis_batch_keys(self):
|
|
"""Test batch key generation with tenant prefix"""
|
|
keys = ["user:1", "user:2", "session:abc"]
|
|
tenant_keys = RedisTenantHelper.batch_make_tenant_keys("tenant-a", "kb-1", keys)
|
|
|
|
assert len(tenant_keys) == 3
|
|
assert all(k.startswith("tenant-a:kb-1:") for k in tenant_keys)
|
|
assert tenant_keys[0] == "tenant-a:kb-1:user:1"
|
|
assert tenant_keys[2] == "tenant-a:kb-1:session:abc"
|
|
|
|
def test_redis_keys_no_collision(self):
|
|
"""Test that tenant/KB combinations create isolated namespaces"""
|
|
key_a = RedisTenantHelper.make_tenant_key("tenant-a", "kb-1", "cache:key")
|
|
key_b = RedisTenantHelper.make_tenant_key("tenant-b", "kb-1", "cache:key")
|
|
|
|
assert key_a != key_b
|
|
assert key_a == "tenant-a:kb-1:cache:key"
|
|
assert key_b == "tenant-b:kb-1:cache:key"
|
|
|
|
|
|
# ============================================================================
|
|
# Context Propagation Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestContextPropagation:
|
|
"""Test tenant context propagation through request pipeline"""
|
|
|
|
def test_tenant_context_creation(self):
|
|
"""Test creating tenant context"""
|
|
context = TenantContext(
|
|
tenant_id="tenant-a", kb_id="kb-1", user_id="user-123", role="admin"
|
|
)
|
|
|
|
assert context.tenant_id == "tenant-a"
|
|
assert context.kb_id == "kb-1"
|
|
assert context.user_id == "user-123"
|
|
assert context.role == "admin"
|
|
|
|
def test_tenant_context_default_values(self):
|
|
"""Test tenant context with minimal data"""
|
|
context = TenantContext(
|
|
tenant_id="default", kb_id="default", user_id="user-default", role="viewer"
|
|
)
|
|
|
|
assert context.tenant_id == "default"
|
|
assert context.kb_id == "default"
|
|
assert context.user_id == "user-default"
|
|
assert context.role == "viewer"
|
|
|
|
|
|
# ============================================================================
|
|
# Tenant Management Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestTenantManagement:
|
|
"""Test tenant CRUD operations"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_tenant(self, mock_tenant_service, sample_tenants):
|
|
"""Test creating a new tenant"""
|
|
tenant_data = sample_tenants["tenant_a"]
|
|
|
|
# Mock the service
|
|
mock_tenant_service.create_tenant = AsyncMock(
|
|
return_value=Tenant(
|
|
tenant_id=tenant_data["tenant_id"],
|
|
tenant_name=tenant_data["name"],
|
|
description=tenant_data["description"],
|
|
)
|
|
)
|
|
|
|
tenant = await mock_tenant_service.create_tenant(
|
|
tenant_name=tenant_data["name"], description=tenant_data["description"]
|
|
)
|
|
|
|
assert tenant.tenant_id == tenant_data["tenant_id"]
|
|
assert tenant.tenant_name == tenant_data["name"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tenants(self, mock_tenant_service, sample_tenants):
|
|
"""Test listing all tenants"""
|
|
tenants_data = list(sample_tenants.values())
|
|
|
|
mock_tenant_service.list_tenants = AsyncMock(
|
|
return_value={
|
|
"items": [
|
|
Tenant(
|
|
tenant_id=t["tenant_id"],
|
|
tenant_name=t["name"],
|
|
description=t["description"],
|
|
)
|
|
for t in tenants_data
|
|
],
|
|
"total": len(tenants_data),
|
|
}
|
|
)
|
|
|
|
result = await mock_tenant_service.list_tenants()
|
|
|
|
assert len(result["items"]) == 2
|
|
assert result["total"] == 2
|
|
|
|
|
|
# ============================================================================
|
|
# Knowledge Base Management Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestKnowledgeBaseManagement:
|
|
"""Test KB CRUD operations with tenant isolation"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_kb_tenant_isolation(self, mock_tenant_service, sample_kbs):
|
|
"""Test that KBs are isolated by tenant"""
|
|
kb_a_list = [kb for kb in sample_kbs.values() if kb["tenant_id"] == "tenant-a"]
|
|
kb_b_list = [kb for kb in sample_kbs.values() if kb["tenant_id"] == "tenant-b"]
|
|
|
|
assert len(kb_a_list) == 2
|
|
assert len(kb_b_list) == 1
|
|
|
|
# Verify no KB ID collision
|
|
kb_a_ids = set(kb["kb_id"] for kb in kb_a_list)
|
|
kb_b_ids = set(kb["kb_id"] for kb in kb_b_list)
|
|
assert len(kb_a_ids & kb_b_ids) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_kb_for_tenant(self, mock_tenant_service, sample_kbs):
|
|
"""Test creating KB within tenant scope"""
|
|
kb_data = sample_kbs["kb_a1"]
|
|
|
|
mock_tenant_service.create_knowledge_base = AsyncMock(
|
|
return_value=KnowledgeBase(
|
|
kb_id=kb_data["kb_id"],
|
|
tenant_id=kb_data["tenant_id"],
|
|
kb_name=kb_data["name"],
|
|
description=kb_data["description"],
|
|
)
|
|
)
|
|
|
|
kb = await mock_tenant_service.create_knowledge_base(
|
|
tenant_id=kb_data["tenant_id"],
|
|
kb_name=kb_data["name"],
|
|
description=kb_data["description"],
|
|
)
|
|
|
|
assert kb.tenant_id == kb_data["tenant_id"]
|
|
assert kb.kb_id == kb_data["kb_id"]
|
|
|
|
|
|
# ============================================================================
|
|
# Document Operation Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestDocumentOperations:
|
|
"""Test document CRUD with tenant/KB isolation"""
|
|
|
|
def test_document_query_by_tenant_kb(self, sample_documents):
|
|
"""Test querying documents scoped to tenant and KB"""
|
|
# Query for tenant-a, kb-a-1
|
|
query_results = [
|
|
d
|
|
for d in sample_documents.values()
|
|
if d["tenant_id"] == "tenant-a" and d["kb_id"] == "kb-a-1"
|
|
]
|
|
|
|
assert len(query_results) == 2
|
|
assert all(d["tenant_id"] == "tenant-a" for d in query_results)
|
|
assert all(d["kb_id"] == "kb-a-1" for d in query_results)
|
|
|
|
def test_cross_tenant_document_access_prevention(self, sample_documents):
|
|
"""Test that cross-tenant document access is prevented"""
|
|
# Try to access tenant-b documents as tenant-a
|
|
tenant_a_docs = [
|
|
d for d in sample_documents.values() if d["tenant_id"] == "tenant-a"
|
|
]
|
|
tenant_b_docs = [
|
|
d for d in sample_documents.values() if d["tenant_id"] == "tenant-b"
|
|
]
|
|
|
|
# Should have no overlap
|
|
tenant_a_ids = set(d["doc_id"] for d in tenant_a_docs)
|
|
tenant_b_ids = set(d["doc_id"] for d in tenant_b_docs)
|
|
|
|
assert len(tenant_a_ids & tenant_b_ids) == 0
|
|
|
|
def test_document_status_isolation(self, sample_documents):
|
|
"""Test that document status is tracked per tenant/KB"""
|
|
docs_a1 = [
|
|
d
|
|
for d in sample_documents.values()
|
|
if d["tenant_id"] == "tenant-a" and d["kb_id"] == "kb-a-1"
|
|
]
|
|
|
|
# All docs should have same status
|
|
assert all(d["status"] == "active" for d in docs_a1)
|
|
|
|
# Create a new doc for a different KB and verify status is independent
|
|
docs_b1 = [
|
|
d
|
|
for d in sample_documents.values()
|
|
if d["tenant_id"] == "tenant-b" and d["kb_id"] == "kb-b-1"
|
|
]
|
|
|
|
# Status can be different
|
|
assert docs_a1[0]["status"] == docs_b1[0]["status"] # Same in this case
|
|
assert docs_a1[0]["tenant_id"] != docs_b1[0]["tenant_id"]
|
|
|
|
|
|
# ============================================================================
|
|
# Entity & Relation Isolation Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestEntityRelationIsolation:
|
|
"""Test entity and relation isolation in graph storage"""
|
|
|
|
def test_entity_tenant_isolation(self):
|
|
"""Test entities are isolated by tenant"""
|
|
entities_a = [
|
|
{
|
|
"entity_id": "e1",
|
|
"tenant_id": "tenant-a",
|
|
"kb_id": "kb-a-1",
|
|
"name": "Entity A1",
|
|
},
|
|
{
|
|
"entity_id": "e2",
|
|
"tenant_id": "tenant-a",
|
|
"kb_id": "kb-a-1",
|
|
"name": "Entity A2",
|
|
},
|
|
]
|
|
entities_b = [
|
|
{
|
|
"entity_id": "e1",
|
|
"tenant_id": "tenant-b",
|
|
"kb_id": "kb-b-1",
|
|
"name": "Entity B1",
|
|
},
|
|
]
|
|
|
|
# Same entity ID but different tenant should be different entities
|
|
entity_a_e1 = next(e for e in entities_a if e["entity_id"] == "e1")
|
|
entity_b_e1 = next(e for e in entities_b if e["entity_id"] == "e1")
|
|
|
|
assert entity_a_e1["tenant_id"] != entity_b_e1["tenant_id"]
|
|
assert entity_a_e1["name"] != entity_b_e1["name"]
|
|
|
|
def test_relation_tenant_isolation(self):
|
|
"""Test relations are isolated by tenant"""
|
|
relations_a = [
|
|
{
|
|
"rel_id": "r1",
|
|
"tenant_id": "tenant-a",
|
|
"kb_id": "kb-a-1",
|
|
"source": "e1",
|
|
"target": "e2",
|
|
"type": "relates_to",
|
|
}
|
|
]
|
|
relations_b = [
|
|
{
|
|
"rel_id": "r1",
|
|
"tenant_id": "tenant-b",
|
|
"kb_id": "kb-b-1",
|
|
"source": "e1",
|
|
"target": "e3",
|
|
"type": "belongs_to",
|
|
}
|
|
]
|
|
|
|
# Same rel_id but different tenant
|
|
rel_a = relations_a[0]
|
|
rel_b = relations_b[0]
|
|
|
|
assert rel_a["tenant_id"] != rel_b["tenant_id"]
|
|
assert rel_a["type"] != rel_b["type"]
|
|
|
|
|
|
# ============================================================================
|
|
# Edge Cases & Error Handling
|
|
# ============================================================================
|
|
|
|
|
|
class TestEdgeCases:
|
|
"""Test edge cases and error conditions"""
|
|
|
|
def test_empty_tenant_id(self):
|
|
"""Test handling of empty tenant ID"""
|
|
# Empty or None tenant ID might not raise error in get_composite_key
|
|
# (depends on implementation), so we test that key is created
|
|
try:
|
|
key = get_composite_key("", "kb-1", "doc-1")
|
|
# If it doesn't raise, key should at least be a string
|
|
assert isinstance(key, str)
|
|
except (ValueError, TypeError, AssertionError):
|
|
# Expected behavior in some implementations
|
|
pass
|
|
|
|
def test_empty_kb_id(self):
|
|
"""Test handling of empty KB ID"""
|
|
try:
|
|
key = get_composite_key("tenant-a", "", "doc-1")
|
|
assert isinstance(key, str)
|
|
except (ValueError, TypeError, AssertionError):
|
|
# Expected behavior in some implementations
|
|
pass
|
|
|
|
def test_composite_key_with_colons(self):
|
|
"""Test composite key generation when parts contain colons (escaped)"""
|
|
# If parts contain colons, they should be handled safely
|
|
key = get_composite_key("tenant:a", "kb:1", "doc:1")
|
|
# Should still create valid composite key
|
|
assert isinstance(key, str)
|
|
assert len(key) > 0
|
|
|
|
def test_very_long_ids(self):
|
|
"""Test composite key with very long IDs"""
|
|
long_tenant_id = "tenant-" + "a" * 1000
|
|
long_kb_id = "kb-" + "b" * 1000
|
|
long_doc_id = "doc-" + "c" * 1000
|
|
|
|
key = get_composite_key(long_tenant_id, long_kb_id, long_doc_id)
|
|
assert isinstance(key, str)
|
|
assert long_tenant_id in key
|
|
assert long_kb_id in key
|
|
assert long_doc_id in key
|
|
|
|
def test_unicode_tenant_ids(self):
|
|
"""Test composite keys with unicode characters"""
|
|
key = get_composite_key("テナント", "知識ベース", "ドキュメント")
|
|
assert isinstance(key, str)
|
|
assert len(key) > 0
|
|
|
|
|
|
# ============================================================================
|
|
# Concurrent Access Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestConcurrentAccess:
|
|
"""Test concurrent multi-tenant operations"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_document_queries(self, sample_documents):
|
|
"""Test concurrent queries from different tenants"""
|
|
|
|
async def query_tenant_docs(tenant_id):
|
|
docs = [d for d in sample_documents.values() if d["tenant_id"] == tenant_id]
|
|
await asyncio.sleep(0.01) # Simulate async operation
|
|
return docs
|
|
|
|
# Query both tenants concurrently
|
|
results = await asyncio.gather(
|
|
query_tenant_docs("tenant-a"), query_tenant_docs("tenant-b")
|
|
)
|
|
|
|
tenant_a_docs = results[0]
|
|
tenant_b_docs = results[1]
|
|
|
|
assert len(tenant_a_docs) == 3
|
|
assert len(tenant_b_docs) == 1
|
|
assert all(d["tenant_id"] == "tenant-a" for d in tenant_a_docs)
|
|
assert all(d["tenant_id"] == "tenant-b" for d in tenant_b_docs)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_kb_operations(self, sample_kbs):
|
|
"""Test concurrent KB operations across tenants"""
|
|
|
|
async def get_tenant_kbs(tenant_id):
|
|
kbs = [kb for kb in sample_kbs.values() if kb["tenant_id"] == tenant_id]
|
|
await asyncio.sleep(0.01)
|
|
return kbs
|
|
|
|
results = await asyncio.gather(
|
|
get_tenant_kbs("tenant-a"), get_tenant_kbs("tenant-b")
|
|
)
|
|
|
|
assert len(results[0]) == 2
|
|
assert len(results[1]) == 1
|
|
|
|
|
|
# ============================================================================
|
|
# Data Consistency Tests
|
|
# ============================================================================
|
|
|
|
|
|
class TestDataConsistency:
|
|
"""Test data consistency across operations"""
|
|
|
|
def test_document_count_by_tenant(self, sample_documents):
|
|
"""Test accurate document counting per tenant"""
|
|
tenant_a_count = len(
|
|
[d for d in sample_documents.values() if d["tenant_id"] == "tenant-a"]
|
|
)
|
|
tenant_b_count = len(
|
|
[d for d in sample_documents.values() if d["tenant_id"] == "tenant-b"]
|
|
)
|
|
|
|
assert tenant_a_count == 3
|
|
assert tenant_b_count == 1
|
|
assert tenant_a_count + tenant_b_count == len(sample_documents)
|
|
|
|
def test_kb_document_consistency(self, sample_documents):
|
|
"""Test document-KB relationships are consistent"""
|
|
kb_a1_docs = [d for d in sample_documents.values() if d["kb_id"] == "kb-a-1"]
|
|
kb_a2_docs = [d for d in sample_documents.values() if d["kb_id"] == "kb-a-2"]
|
|
|
|
# Verify consistency: all docs in KB should have matching tenant
|
|
assert all(d["tenant_id"] == "tenant-a" for d in kb_a1_docs)
|
|
assert all(d["tenant_id"] == "tenant-a" for d in kb_a2_docs)
|
|
|
|
# Verify no doc appears in multiple KBs
|
|
kb_a1_ids = set(d["doc_id"] for d in kb_a1_docs)
|
|
kb_a2_ids = set(d["doc_id"] for d in kb_a2_docs)
|
|
assert len(kb_a1_ids & kb_a2_ids) == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Run tests
|
|
pytest.main([__file__, "-v", "--tb=short"])
|