206 lines
6.8 KiB
Python
206 lines
6.8 KiB
Python
"""
|
|
Tests for tenant isolation in multi-tenant LightRAG.
|
|
|
|
These tests verify that:
|
|
1. Tenants can be created and managed
|
|
2. Knowledge bases are properly isolated per tenant
|
|
3. Cross-tenant data access is denied
|
|
4. Storage operations respect tenant/KB boundaries
|
|
"""
|
|
|
|
import pytest
|
|
from lightrag.models.tenant import (
|
|
Tenant,
|
|
TenantConfig,
|
|
TenantContext,
|
|
KnowledgeBase,
|
|
KBConfig,
|
|
ResourceQuota,
|
|
)
|
|
|
|
|
|
class TestTenantModels:
|
|
"""Test tenant data models."""
|
|
|
|
def test_tenant_creation(self):
|
|
"""Test creating a tenant with default values."""
|
|
tenant = Tenant(tenant_name="Test Tenant")
|
|
assert tenant.tenant_name == "Test Tenant"
|
|
assert tenant.is_active is True
|
|
assert tenant.tenant_id is not None
|
|
assert len(tenant.tenant_id) > 0
|
|
assert tenant.kb_count == 0
|
|
|
|
def test_tenant_with_custom_config(self):
|
|
"""Test creating a tenant with custom configuration."""
|
|
config = TenantConfig(
|
|
llm_model="custom-llm",
|
|
embedding_model="custom-embedding",
|
|
chunk_size=2000,
|
|
top_k=50,
|
|
)
|
|
tenant = Tenant(tenant_name="Test", config=config)
|
|
assert tenant.config.llm_model == "custom-llm"
|
|
assert tenant.config.chunk_size == 2000
|
|
assert tenant.config.top_k == 50
|
|
|
|
def test_tenant_to_dict(self):
|
|
"""Test tenant serialization to dictionary."""
|
|
tenant = Tenant(
|
|
tenant_id="test-123", tenant_name="Test", description="Test tenant"
|
|
)
|
|
data = tenant.to_dict()
|
|
assert data["tenant_id"] == "test-123"
|
|
assert data["tenant_name"] == "Test"
|
|
assert data["description"] == "Test tenant"
|
|
assert "config" in data
|
|
assert "quota" in data
|
|
assert data["is_active"] is True
|
|
|
|
|
|
class TestKnowledgeBase:
|
|
"""Test knowledge base models."""
|
|
|
|
def test_kb_creation(self):
|
|
"""Test creating a knowledge base."""
|
|
kb = KnowledgeBase(tenant_id="tenant-1", kb_name="Test KB")
|
|
assert kb.kb_name == "Test KB"
|
|
assert kb.tenant_id == "tenant-1"
|
|
assert kb.is_active is True
|
|
assert kb.status == "ready"
|
|
assert kb.document_count == 0
|
|
|
|
def test_kb_to_dict(self):
|
|
"""Test KB serialization."""
|
|
kb = KnowledgeBase(kb_id="kb-123", tenant_id="tenant-1", kb_name="Test KB")
|
|
data = kb.to_dict()
|
|
assert data["kb_id"] == "kb-123"
|
|
assert data["tenant_id"] == "tenant-1"
|
|
assert data["kb_name"] == "Test KB"
|
|
|
|
def test_kb_with_override_config(self):
|
|
"""Test KB with configuration override."""
|
|
config = KBConfig(top_k=30, chunk_size=1500)
|
|
kb = KnowledgeBase(tenant_id="tenant-1", kb_name="Test", config=config)
|
|
assert kb.config.top_k == 30
|
|
assert kb.config.chunk_size == 1500
|
|
|
|
|
|
class TestTenantContext:
|
|
"""Test tenant context for requests."""
|
|
|
|
def test_context_creation(self):
|
|
"""Test creating a tenant context."""
|
|
context = TenantContext(
|
|
tenant_id="tenant-1", kb_id="kb-1", user_id="user-1", role="admin"
|
|
)
|
|
assert context.tenant_id == "tenant-1"
|
|
assert context.kb_id == "kb-1"
|
|
assert context.user_id == "user-1"
|
|
assert context.role == "admin"
|
|
|
|
def test_workspace_namespace(self):
|
|
"""Test backward-compatible workspace namespace."""
|
|
context = TenantContext(
|
|
tenant_id="acme", kb_id="docs", user_id="user-1", role="editor"
|
|
)
|
|
assert context.workspace_namespace == "acme_docs"
|
|
|
|
def test_kb_access_control(self):
|
|
"""Test KB access control in context."""
|
|
context = TenantContext(
|
|
tenant_id="tenant-1",
|
|
kb_id="kb-1",
|
|
user_id="user-1",
|
|
role="viewer",
|
|
knowledge_base_ids=["kb-1", "kb-2"],
|
|
)
|
|
assert context.can_access_kb("kb-1") is True
|
|
assert context.can_access_kb("kb-2") is True
|
|
assert context.can_access_kb("kb-3") is False
|
|
|
|
def test_kb_access_all(self):
|
|
"""Test KB access with wildcard."""
|
|
context = TenantContext(
|
|
tenant_id="tenant-1",
|
|
kb_id="kb-1",
|
|
user_id="user-1",
|
|
role="admin",
|
|
knowledge_base_ids=["*"], # Admin has access to all
|
|
)
|
|
assert context.can_access_kb("kb-1") is True
|
|
assert context.can_access_kb("kb-2") is True
|
|
assert context.can_access_kb("kb-999") is True
|
|
|
|
def test_permission_checking(self):
|
|
"""Test permission checking."""
|
|
context = TenantContext(
|
|
tenant_id="tenant-1",
|
|
kb_id="kb-1",
|
|
user_id="user-1",
|
|
role="viewer",
|
|
permissions={"query:run": True, "document:create": False},
|
|
)
|
|
assert context.has_permission("query:run") is True
|
|
assert context.has_permission("document:create") is False
|
|
|
|
def test_context_to_dict(self):
|
|
"""Test context serialization."""
|
|
context = TenantContext(
|
|
tenant_id="tenant-1", kb_id="kb-1", user_id="user-1", role="editor"
|
|
)
|
|
data = context.to_dict()
|
|
assert data["tenant_id"] == "tenant-1"
|
|
assert data["kb_id"] == "kb-1"
|
|
assert data["user_id"] == "user-1"
|
|
assert data["workspace_namespace"] == "tenant-1_kb-1"
|
|
|
|
|
|
class TestResourceQuota:
|
|
"""Test resource quota limits."""
|
|
|
|
def test_quota_defaults(self):
|
|
"""Test quota has reasonable defaults."""
|
|
quota = ResourceQuota()
|
|
assert quota.max_documents == 10000
|
|
assert quota.max_storage_gb == 100.0
|
|
assert quota.max_concurrent_queries == 10
|
|
assert quota.max_monthly_api_calls == 100000
|
|
|
|
def test_quota_custom(self):
|
|
"""Test custom quota settings."""
|
|
quota = ResourceQuota(
|
|
max_documents=5000, max_storage_gb=50.0, max_concurrent_queries=5
|
|
)
|
|
assert quota.max_documents == 5000
|
|
assert quota.max_storage_gb == 50.0
|
|
assert quota.max_concurrent_queries == 5
|
|
|
|
|
|
class TestRoles:
|
|
"""Test role and permission definitions."""
|
|
|
|
def test_role_enum_values(self):
|
|
"""Test role enum has expected values."""
|
|
from lightrag.models.tenant import Role
|
|
|
|
assert Role.ADMIN.value == "admin"
|
|
assert Role.EDITOR.value == "editor"
|
|
assert Role.VIEWER.value == "viewer"
|
|
assert Role.VIEWER_READONLY.value == "viewer:read-only"
|
|
|
|
def test_role_permissions(self):
|
|
"""Test role-to-permissions mapping."""
|
|
from lightrag.models.tenant import ROLE_PERMISSIONS, Role
|
|
|
|
admin_perms = ROLE_PERMISSIONS[Role.ADMIN]
|
|
assert len(admin_perms) > 0
|
|
assert "query:run" in admin_perms
|
|
|
|
viewer_perms = ROLE_PERMISSIONS[Role.VIEWER]
|
|
assert "query:run" in viewer_perms
|
|
assert "document:delete" not in viewer_perms
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|