LightRAG/tests/test_document_routes_tenant_scoped.py

"""
Integration tests for tenant-scoped document routes (Phase 3).

This test suite verifies that all document-related endpoints properly use
tenant-scoped RAG instances, ensuring multi-tenant data isolation and correct
document visibility within each tenant's knowledge base.

Key test scenarios:
1. Document text insertion via /text endpoint (tenant-scoped)
2. Batch text insertion via /texts endpoint (tenant-scoped)
3. Document listing via /documents endpoint (tenant-scoped)
4. Document status tracking via /track_status endpoint (tenant-scoped)
5. Multi-tenant isolation: documents in Tenant A are not visible in Tenant B
"""

import pytest
from datetime import datetime
from unittest.mock import AsyncMock
from fastapi.testclient import TestClient
from fastapi import FastAPI, Depends

from lightrag.models.tenant import TenantContext, Role
from lightrag.base import DocProcessingStatus, DocStatus
from lightrag.api.dependencies import get_tenant_context
from lightrag.api.routers.document_routes import create_document_routes
from lightrag.lightrag import LightRAG


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

@pytest.fixture
def tenant_context_a():
    """Tenant A context for isolation testing."""
    return TenantContext(
        tenant_id="tenant-a",
        kb_id="kb-a-1",
        user_id="user-a-123",
        role=Role.ADMIN
    )


@pytest.fixture
def tenant_context_b():
    """Tenant B context for isolation testing."""
    return TenantContext(
        tenant_id="tenant-b",
        kb_id="kb-b-1",
        user_id="user-b-456",
        role=Role.ADMIN
    )


@pytest.fixture
def mock_rag_instance():
    """Create a mock LightRAG instance for testing."""
    mock = AsyncMock(spec=LightRAG)
    
    # Mock doc_status storage
    mock.doc_status = AsyncMock()
    mock.doc_status.get_doc_by_file_path = AsyncMock(return_value=None)
    mock.doc_status.get_doc = AsyncMock(return_value=None)
    
    # Mock document retrieval methods
    mock.get_docs_by_status = AsyncMock(return_value={})
    mock.aget_docs_by_track_id = AsyncMock(return_value={})
    
    return mock


@pytest.fixture
def mock_rag_instances():
    """Create separate mock RAG instances for each tenant for isolation testing."""
    rag_a = AsyncMock(spec=LightRAG)
    rag_b = AsyncMock(spec=LightRAG)
    
    # Setup Tenant A RAG
    rag_a.doc_status = AsyncMock()
    rag_a.doc_status.get_doc_by_file_path = AsyncMock(return_value=None)
    rag_a.get_docs_by_status = AsyncMock(return_value={
        "doc-a-1": DocProcessingStatus(
            content_summary="Sample doc in Tenant A",
            content_length=100,
            file_path="doc-a-1",
            status=DocStatus.PROCESSED,
            created_at=datetime.utcnow().isoformat(),
            updated_at=datetime.utcnow().isoformat(),
            track_id="track-a-1",
        )
    })
    rag_a.aget_docs_by_track_id = AsyncMock(return_value={
        "doc-a-1": DocProcessingStatus(
            content_summary="Sample doc in Tenant A",
            content_length=100,
            file_path="doc-a-1",
            status=DocStatus.PROCESSED,
            created_at=datetime.utcnow().isoformat(),
            updated_at=datetime.utcnow().isoformat(),
            track_id="track-a-1",
        )
    })
    
    # Setup Tenant B RAG
    rag_b.doc_status = AsyncMock()
    rag_b.doc_status.get_doc_by_file_path = AsyncMock(return_value=None)
    rag_b.get_docs_by_status = AsyncMock(return_value={
        "doc-b-1": DocProcessingStatus(
            content_summary="Sample doc in Tenant B",
            content_length=200,
            file_path="doc-b-1",
            status=DocStatus.PROCESSED,
            created_at=datetime.utcnow().isoformat(),
            updated_at=datetime.utcnow().isoformat(),
            track_id="track-b-1",
        )
    })
    rag_b.aget_docs_by_track_id = AsyncMock(return_value={
        "doc-b-1": DocProcessingStatus(
            content_summary="Sample doc in Tenant B",
            content_length=200,
            file_path="doc-b-1",
            status=DocStatus.PROCESSED,
            created_at=datetime.utcnow().isoformat(),
            updated_at=datetime.utcnow().isoformat(),
            track_id="track-b-1",
        )
    })
    
    return {"tenant-a": rag_a, "tenant-b": rag_b}


@pytest.fixture
def app_with_document_routes(mock_rag_instances):
    """Create FastAPI app with document routes and mocked dependencies."""
    app = FastAPI()

    # Build router with a custom rag_manager so tests can provide mock rag instances
    # Create a tiny rag_manager that returns the correct mock for a tenant id
    class DummyRagManager:
        def __init__(self, mapping):
            self.mapping = mapping

        async def get_rag_instance(self, tenant_id, kb_id, user_id=None):
            return self.mapping.get(tenant_id)

    dummy_rag = AsyncMock(spec=LightRAG)
    dummy_doc_manager = AsyncMock()
    rag_manager = DummyRagManager(mock_rag_instances)

    # Create and include a router instance pointing to our rag_manager
    doc_router = create_document_routes(dummy_rag, dummy_doc_manager, rag_manager=rag_manager)
    # The document router already defines prefix="/documents" internally.
    # Mount it under "/api" so final paths become "/api/documents/...".
    app.include_router(doc_router, prefix="/api", tags=["documents"])
    
    # Track the current tenant context
    current_context = {"context": None, "rag": None}
    
    # Mock get_tenant_context
    async def mock_get_tenant_context(*args, **kwargs):
        return current_context["context"]
    
    # Override dependencies
    # The router uses get_tenant_context_optional in its internal dependency chain
    # so override both get_tenant_context and get_tenant_context_optional to ensure
    # tenant context is injected in tests (strict multi-tenant mode requires it).
    from lightrag.api.dependencies import get_tenant_context as _gtc
    from lightrag.api.dependencies import get_tenant_context_optional as _gtco

    app.dependency_overrides[_gtc] = mock_get_tenant_context
    app.dependency_overrides[_gtco] = mock_get_tenant_context

    # For testing route logic we don't want auth to interfere. Disable global auth checks
    # by toggling the module-level flag in utils_api. This mirrors the behavior in
    # many test fixtures where authentication is mocked out.
    import lightrag.api.utils_api as utils_api
    utils_api.auth_configured = False
    
    # Store context setter for tests
    app._set_context = lambda ctx: current_context.update({"context": ctx})
    
    return app


@pytest.fixture
def client(app_with_document_routes):
    """Create test client."""
    return TestClient(app_with_document_routes)


# ============================================================================
# Test Cases
# ============================================================================

class TestDocumentRoutesUseTenantRAG:
    """Test that document routes properly use tenant-scoped RAG instances."""
    
    def test_text_endpoint_uses_tenant_rag(self, client, app_with_document_routes, tenant_context_a, mock_rag_instances):
        """Test that /text endpoint uses tenant-specific RAG instance."""
        app_with_document_routes._set_context(tenant_context_a)
        
        response = client.post(
            "/api/documents/text?args=1&kwargs=1",
            json={
                "text": "This is a test document",
                "file_source": "test_file.txt"
            }
        )
        
        # Verify that the request was processed
        if response.status_code == 422:
            # Debugging info: show validation errors
            print("RESPONSE 422 BODY:", response.json())
        assert response.status_code in [200, 400, 500], f"Unexpected status code: {response.status_code}"
        
        # Verify that tenant A's RAG was queried for doc status
        mock_rag_instances["tenant-a"].doc_status.get_doc_by_file_path.assert_called()
    
    def test_texts_endpoint_uses_tenant_rag(self, client, app_with_document_routes, tenant_context_a, mock_rag_instances):
        """Test that /texts endpoint uses tenant-specific RAG instance."""
        app_with_document_routes._set_context(tenant_context_a)
        
        response = client.post(
            "/api/documents/texts?args=1&kwargs=1",
            json={
                "texts": ["Document 1", "Document 2"],
                "file_sources": ["doc1.txt", "doc2.txt"]
            }
        )
        
        # Verify that the request was processed
        if response.status_code == 422:
            print("RESPONSE 422 BODY:", response.json())
        assert response.status_code in [200, 400, 500]
        
        # Verify that tenant A's RAG was queried for doc status
        mock_rag_instances["tenant-a"].doc_status.get_doc_by_file_path.assert_called()
    
    def test_documents_endpoint_uses_tenant_rag(self, client, app_with_document_routes, tenant_context_a, mock_rag_instances):
        """Test that /documents GET endpoint uses tenant-specific RAG instance."""
        app_with_document_routes._set_context(tenant_context_a)
        
        response = client.get("/api/documents?args=1&kwargs=1")
        
        # Verify that the request was processed
        if response.status_code == 422:
            print("RESPONSE 422 BODY:", response.json())
        assert response.status_code == 200
        
        # Verify that tenant A's RAG was used to get docs by status
        mock_rag_instances["tenant-a"].get_docs_by_status.assert_called()
    
    def test_track_status_endpoint_uses_tenant_rag(self, client, app_with_document_routes, tenant_context_a, mock_rag_instances):
        """Test that /track_status endpoint uses tenant-specific RAG instance."""
        app_with_document_routes._set_context(tenant_context_a)
        
        response = client.get("/api/documents/track_status/track-a-1?args=1&kwargs=1")
        
        # Verify that the request was processed
        if response.status_code == 422:
            print("RESPONSE 422 BODY:", response.json())
        assert response.status_code == 200
        
        # Verify that tenant A's RAG was used to get docs by track_id
        mock_rag_instances["tenant-a"].aget_docs_by_track_id.assert_called()


class TestMultiTenantIsolation:
    """Test that documents from different tenants are isolated from each other."""
    
    def test_tenant_a_cannot_see_tenant_b_documents(self, client, app_with_document_routes, tenant_context_a, tenant_context_b, mock_rag_instances):
        """Verify that Tenant A's document queries don't return Tenant B's documents."""
        # Query as Tenant A
        app_with_document_routes._set_context(tenant_context_a)
        response_a = client.get("/api/documents?args=1&kwargs=1")
        
        # Query as Tenant B
        app_with_document_routes._set_context(tenant_context_b)
        response_b = client.get("/api/documents?args=1&kwargs=1")
        
        # Both should succeed
        assert response_a.status_code == 200
        assert response_b.status_code == 200
        
        # Verify that different RAG instances were used
        mock_rag_instances["tenant-a"].get_docs_by_status.assert_called()
        mock_rag_instances["tenant-b"].get_docs_by_status.assert_called()
    
    def test_track_status_returns_only_tenant_documents(self, client, app_with_document_routes, tenant_context_a, tenant_context_b, mock_rag_instances):
        """Verify that track_status endpoint returns docs from the correct tenant only."""
        # Track status in Tenant A
        app_with_document_routes._set_context(tenant_context_a)
        response_a = client.get("/api/documents/track_status/track-a-1?args=1&kwargs=1")
        
        # Track status in Tenant B
        app_with_document_routes._set_context(tenant_context_b)
        response_b = client.get("/api/documents/track_status/track-b-1?args=1&kwargs=1")
        
        # Both should succeed
        assert response_a.status_code == 200
        assert response_b.status_code == 200
        
        # Verify that different RAG instances were queried for different track IDs
        mock_rag_instances["tenant-a"].aget_docs_by_track_id.assert_called_with("track-a-1")
        mock_rag_instances["tenant-b"].aget_docs_by_track_id.assert_called_with("track-b-1")


class TestDocumentEndpointFunctionality:
    """Test basic functionality of document endpoints."""
    
    def test_text_endpoint_duplicate_file_source_rejection(self, client, app_with_document_routes, tenant_context_a, mock_rag_instances):
        """Test that /text endpoint rejects duplicate file sources."""
        app_with_document_routes._set_context(tenant_context_a)
        
        # Mock that file already exists
        existing_doc = {"status": "PROCESSED"}
        mock_rag_instances["tenant-a"].doc_status.get_doc_by_file_path.return_value = existing_doc
        
        response = client.post(
            "/api/documents/text?args=1&kwargs=1",
            json={
                "text": "Duplicate content",
                "file_source": "duplicate.txt"
            }
        )
        
        # Should return duplicated status
        assert response.status_code == 200
        data = response.json()
        assert data.get("status") == "duplicated"
    
    def test_texts_endpoint_accepts_batch_insert(self, client, app_with_document_routes, tenant_context_a, mock_rag_instances):
        """Test that /texts endpoint can handle batch text insertion."""
        app_with_document_routes._set_context(tenant_context_a)
        
        response = client.post(
            "/api/documents/texts?args=1&kwargs=1",
            json={
                "texts": ["Text 1", "Text 2", "Text 3"],
                "file_sources": ["file1.txt", "file2.txt", "file3.txt"]
            }
        )
        
        # Should accept the batch
        assert response.status_code == 200
        data = response.json()
        assert data.get("status") == "success"


# ============================================================================
# Test Runner Configuration
# ============================================================================

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