LightRAG/tests/test_tenant_storage_phase3.py

"""
Phase 3 Unit Tests: Tenant Storage Layer Integration

Tests for tenant-aware storage operations, instance caching, and data isolation.
Verifies that the storage layer correctly enforces multi-tenant boundaries.
"""

import pytest
import tempfile
from unittest.mock import MagicMock
from lightrag.tenant_rag_manager import TenantRAGManager
from lightrag.models.tenant import Tenant, TenantContext, TenantConfig
from lightrag.services.tenant_service import TenantService


@pytest.fixture
def temp_working_dir():
    """Create a temporary directory for testing."""
    with tempfile.TemporaryDirectory() as tmpdir:
        yield tmpdir


@pytest.fixture
def mock_tenant_service():
    """Create a mock TenantService."""
    service = MagicMock(spec=TenantService)
    
    async def mock_get_tenant(tenant_id: str):
        if tenant_id.startswith("tenant_"):
            return Tenant(
                tenant_id=tenant_id,
                tenant_name=f"Test {tenant_id}",
                config=TenantConfig(
                    llm_model="gpt-4o-mini",
                    embedding_model="bge-m3:latest",
                    top_k=40,
                    cosine_threshold=0.2,
                ),
                is_active=True
            )
        return None
    
    service.get_tenant = mock_get_tenant
    return service


class TestTenantRAGManagerBasics:
    """Basic unit tests for TenantRAGManager."""
    
    def test_manager_initialization(self, temp_working_dir, mock_tenant_service):
        """Test TenantRAGManager initializes correctly."""
        manager = TenantRAGManager(
            base_working_dir=temp_working_dir,
            tenant_service=mock_tenant_service,
            max_cached_instances=10
        )
        
        assert manager.base_working_dir == temp_working_dir
        assert manager.max_cached_instances == 10
        assert manager.get_instance_count() == 0
    
    def test_manager_str_representation(self, temp_working_dir, mock_tenant_service):
        """Test manager string representation."""
        manager = TenantRAGManager(
            base_working_dir=temp_working_dir,
            tenant_service=mock_tenant_service,
            max_cached_instances=50
        )
        
        repr_str = repr(manager)
        assert "TenantRAGManager" in repr_str
        assert "50" in repr_str
    
    def test_get_cached_keys_empty(self, temp_working_dir, mock_tenant_service):
        """Test get_cached_keys returns empty list when no instances cached."""
        manager = TenantRAGManager(
            base_working_dir=temp_working_dir,
            tenant_service=mock_tenant_service,
        )
        
        keys = manager.get_cached_keys()
        assert isinstance(keys, list)
        assert len(keys) == 0
    
    def test_manager_attributes(self, temp_working_dir, mock_tenant_service):
        """Test manager has expected attributes."""
        manager = TenantRAGManager(
            base_working_dir=temp_working_dir,
            tenant_service=mock_tenant_service,
            max_cached_instances=100
        )
        
        assert hasattr(manager, "_instances")
        assert hasattr(manager, "_lock")
        assert hasattr(manager, "_access_order")
        assert hasattr(manager, "base_working_dir")
        assert hasattr(manager, "tenant_service")
        assert hasattr(manager, "max_cached_instances")


class TestTenantContextUnit:
    """Test suite for TenantContext usage."""
    
    def test_tenant_context_creation(self):
        """Test TenantContext creation."""
        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_tenant_context_with_permissions(self):
        """Test TenantContext with permissions."""
        context = TenantContext(
            tenant_id="tenant_1",
            kb_id="kb_1",
            user_id="user_1",
            role="editor",
            permissions={"query": True, "edit": True, "delete": False}
        )
        
        assert context.permissions["query"] is True
        assert context.permissions["edit"] is True
        assert context.permissions["delete"] is False
    
    def test_tenant_context_workspace_namespace(self):
        """Test workspace namespace generation from TenantContext."""
        context = TenantContext(
            tenant_id="tenant_1",
            kb_id="kb_1",
            user_id="user_1",
            role="admin"
        )
        
        # Should generate workspace namespace for backward compatibility
        workspace = context.workspace_namespace
        assert "tenant_1" in workspace
        assert "kb_1" in workspace
    
    def test_tenant_context_different_roles(self):
        """Test TenantContext with different role values."""
        roles = ["admin", "editor", "viewer", "viewer_readonly"]
        
        for role in roles:
            context = TenantContext(
                tenant_id="tenant_1",
                kb_id="kb_1",
                user_id="user_1",
                role=role
            )
            assert context.role == role
    
    def test_tenant_context_default_permissions(self):
        """Test TenantContext with default permissions."""
        context = TenantContext(
            tenant_id="tenant_1",
            kb_id="kb_1",
            user_id="user_1",
            role="admin"
        )
        
        # Default permissions should be empty dict
        assert isinstance(context.permissions, dict)


class TestTenantModel:
    """Test suite for Tenant data model."""
    
    def test_tenant_creation_basic(self):
        """Test basic Tenant creation."""
        tenant = Tenant(
            tenant_id="test_tenant",
            tenant_name="Test Tenant"
        )
        
        assert tenant.tenant_id == "test_tenant"
        assert tenant.tenant_name == "Test Tenant"
        assert tenant.is_active is True
    
    def test_tenant_creation_with_config(self):
        """Test Tenant creation with custom config."""
        config = TenantConfig(
            llm_model="gpt-4",
            embedding_model="text-embedding-3",
            top_k=50,
        )
        
        tenant = Tenant(
            tenant_id="test_tenant",
            tenant_name="Test Tenant",
            config=config
        )
        
        assert tenant.config.llm_model == "gpt-4"
        assert tenant.config.embedding_model == "text-embedding-3"
        assert tenant.config.top_k == 50
    
    def test_tenant_inactive_status(self):
        """Test Tenant with inactive status."""
        tenant = Tenant(
            tenant_id="test_tenant",
            tenant_name="Inactive Tenant",
            is_active=False
        )
        
        assert tenant.is_active is False
    
    def test_tenant_with_metadata(self):
        """Test Tenant with custom metadata."""
        metadata = {"custom_field": "value", "flag": True}
        tenant = Tenant(
            tenant_id="test_tenant",
            tenant_name="Test",
            metadata=metadata
        )
        
        assert tenant.metadata["custom_field"] == "value"
        assert tenant.metadata["flag"] is True


class TestTenantConfigModel:
    """Test suite for TenantConfig data model."""
    
    def test_config_defaults(self):
        """Test TenantConfig default values."""
        config = TenantConfig()
        
        assert config.llm_model == "gpt-4o-mini"
        assert config.embedding_model == "bge-m3:latest"
        assert config.chunk_size == 1200
        assert config.chunk_overlap == 100
        assert config.top_k == 40
        assert config.cosine_threshold == 0.2
        assert config.enable_llm_cache is True
    
    def test_config_custom_values(self):
        """Test TenantConfig with custom values."""
        config = TenantConfig(
            llm_model="gpt-4-turbo",
            embedding_model="claude-embedding",
            chunk_size=2000,
            top_k=100,
            cosine_threshold=0.5
        )
        
        assert config.llm_model == "gpt-4-turbo"
        assert config.embedding_model == "claude-embedding"
        assert config.chunk_size == 2000
        assert config.top_k == 100
        assert config.cosine_threshold == 0.5
    
    def test_config_custom_metadata(self):
        """Test TenantConfig with custom metadata."""
        custom_meta = {"provider": "custom", "version": "2.0"}
        config = TenantConfig(custom_metadata=custom_meta)
        
        assert config.custom_metadata["provider"] == "custom"
        assert config.custom_metadata["version"] == "2.0"


class TestMultiTenantStructure:
    """Test multi-tenant structure and isolation concepts."""
    
    def test_different_tenant_ids_are_distinct(self):
        """Test that different tenant IDs create distinct objects."""
        tenant1 = Tenant(tenant_id="tenant_1", tenant_name="T1")
        tenant2 = Tenant(tenant_id="tenant_2", tenant_name="T2")
        
        assert tenant1.tenant_id != tenant2.tenant_id
        assert tenant1.tenant_id == "tenant_1"
        assert tenant2.tenant_id == "tenant_2"
    
    def test_tenant_context_scoping(self):
        """Test that TenantContext properly scopes to tenant/kb."""
        context1 = TenantContext(
            tenant_id="t1",
            kb_id="kb1",
            user_id="user1",
            role="admin"
        )
        context2 = TenantContext(
            tenant_id="t2",
            kb_id="kb1",
            user_id="user1",
            role="admin"
        )
        
        # Same KB but different tenants
        assert context1.tenant_id != context2.tenant_id
        assert context1.kb_id == context2.kb_id
        assert context1.workspace_namespace != context2.workspace_namespace
    
    def test_composite_workspace_naming_convention(self):
        """Test composite workspace naming convention."""
        context_t1_kb1 = TenantContext(
            tenant_id="tenant_1",
            kb_id="kb_1",
            user_id="user1",
            role="admin"
        )
        context_t1_kb2 = TenantContext(
            tenant_id="tenant_1",
            kb_id="kb_2",
            user_id="user1",
            role="admin"
        )
        
        ns1 = context_t1_kb1.workspace_namespace
        ns2 = context_t1_kb2.workspace_namespace
        
        # Both should have tenant_1 but different kb_id
        assert "tenant_1" in ns1
        assert "tenant_1" in ns2
        assert "kb_1" in ns1
        assert "kb_2" in ns2
        assert ns1 != ns2
    
    def test_multiple_tenants_same_kb_different_contexts(self):
        """Test multiple tenants with same KB name have different contexts."""
        contexts = []
        for i in range(1, 4):
            ctx = TenantContext(
                tenant_id=f"tenant_{i}",
                kb_id="shared_kb",
                user_id="admin",
                role="admin"
            )
            contexts.append(ctx)
        
        # All should have different workspaces
        namespaces = [c.workspace_namespace for c in contexts]
        assert len(set(namespaces)) == len(namespaces), "Namespaces should be unique"
        
        # All should contain shared_kb
        for ns in namespaces:
            assert "shared_kb" in ns


class TestTenantContextIsolation:
    """Test tenant context isolation and permissions."""
    
    def test_context_prevents_cross_tenant_access_conceptually(self):
        """Test that context structure prevents cross-tenant access."""
        admin_context = TenantContext(
            tenant_id="tenant_1",
            kb_id="kb_1",
            user_id="admin_user",
            role="admin",
            permissions={"query": True, "edit": True, "delete": True}
        )
        
        viewer_context = TenantContext(
            tenant_id="tenant_1",
            kb_id="kb_2",
            user_id="viewer_user",
            role="viewer",
            permissions={"query": True}
        )
        
        # Different KB IDs even in same tenant
        assert admin_context.kb_id != viewer_context.kb_id
        # Different permissions
        assert admin_context.permissions != viewer_context.permissions
    
    def test_context_role_based_access_control(self):
        """Test role-based access control in TenantContext."""
        admin_perms = {"query": True, "edit": True, "delete": True}
        editor_perms = {"query": True, "edit": True, "delete": False}
        viewer_perms = {"query": True, "edit": False, "delete": False}
        
        admin_ctx = TenantContext("t1", "kb1", "u1", "admin", admin_perms)
        editor_ctx = TenantContext("t1", "kb1", "u2", "editor", editor_perms)
        viewer_ctx = TenantContext("t1", "kb1", "u3", "viewer", viewer_perms)
        
        # Admin has all permissions
        assert admin_ctx.permissions["delete"] is True
        # Editor cannot delete
        assert editor_ctx.permissions["delete"] is False
        # Viewer cannot edit
        assert viewer_ctx.permissions["edit"] is False


if __name__ == "__main__":
    pytest.main([__file__, "-v"])