tests/test_pipeline.py

import pytest
from app.core.pipeline import rag_pipeline


@pytest.mark.integration
class TestRAGPipeline:
    
    @pytest.mark.asyncio
    async def test_pipeline_initialization(self):
        assert rag_pipeline.retriever is not None
        assert rag_pipeline.reranker is not None
        assert rag_pipeline.generator is not None
        assert rag_pipeline.cache is not None
        assert rag_pipeline.memory is not None
    
    @pytest.mark.asyncio
    async def test_generate_without_rag(self, clear_cache):
        query = "What is 2 + 2?"
        
        response = await rag_pipeline.generate(
            query=query,
            use_context=False
        )
        
        assert response is not None
        assert isinstance(response, str)
        assert len(response) > 0
    
    @pytest.mark.asyncio
    @pytest.mark.slow
    async def test_generate_with_rag(self, clear_cache):
        query = "Explain the attention mechanism"
        
        try:
            response = await rag_pipeline.generate(
                query=query,
                use_context=True
            )
            
            assert response is not None
            assert isinstance(response, str)
            assert len(response) > 0
        except Exception as e:
            pytest.skip(f"RAG test skipped: {str(e)}")
    
    @pytest.mark.asyncio
    @pytest.mark.slow
    async def test_streaming_response(self, clear_cache):
        query = "What is machine learning?"
        
        chunks = []
        async for chunk in rag_pipeline.stream(
            query=query,
            use_context=False
        ):
            chunks.append(chunk)
        
        assert len(chunks) > 0
        full_response = "".join(chunks)
        assert len(full_response) > 0
    
    @pytest.mark.asyncio
    async def test_pipeline_with_session(self, clear_cache):
        session_id = "test-session-123"
        
        response1 = await rag_pipeline.generate(
            query="My name is Alice",
            session_id=session_id,
            use_context=False
        )
        
        assert response1 is not None
        
        messages = rag_pipeline.memory.get_messages(session_id)
        assert len(messages) >= 2
    
    @pytest.mark.asyncio
    async def test_cache_integration(self, clear_cache):
        query = "What is Python?"
        
        response1 = await rag_pipeline.generate(
            query=query,
            use_context=False
        )
        
        response2 = await rag_pipeline.generate(
            query=query,
            use_context=False
        )
        
        assert response1 == response2