CoolFace
Apppublic

ruby2210/rag-chatbot

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
test_ingestion_pipeline.py148 linesDownload Raw Back to integration
1"""2Integration test for ingestion pipeline in the RAG Chatbot application.3"""4import pytest5import os6from pathlib import Path7from fastapi.testclient import TestClient8from src.api.main import app9from src.services.ingestion_service import ingestion_service10 11 12@pytest.fixture13def client():14    """Create a test client for the API."""15    return TestClient(app)16 17 18def test_ingestion_pipeline_integration(client, tmp_path):19    """Test the complete ingestion pipeline integration."""20    # Create temporary markdown files for testing21    test_docs_dir = tmp_path / "docs"22    test_docs_dir.mkdir()23 24    # Create a test markdown file25    test_file = test_docs_dir / "test_doc.md"26    test_content = """27# Test Document28 29This is a test document for the RAG chatbot.30 31## Section 132 33The concept of RAG (Retrieval-Augmented Generation) is important in modern AI systems.34 35## Section 236 37RAG combines retrieval and generation to provide accurate, context-aware responses.38"""39    test_file.write_text(test_content)40 41    # Prepare ingestion request42    ingest_data = {43        "source_path": str(test_docs_dir)44    }45 46    # Call the ingestion endpoint47    response = client.post("/api/embeddings/ingest", json=ingest_data)48 49    # Verify the response50    assert response.status_code in [200, 500], f"Expected 200 or 500, got {response.status_code}"51 52    if response.status_code == 200:53        data = response.json()54 55        assert "status" in data56        assert "files_processed" in data57        assert "chunks_created" in data58        assert "message" in data59 60        assert data["status"] in ["completed", "failed"]61        assert isinstance(data["files_processed"], int)62        assert isinstance(data["chunks_created"], int)63        assert isinstance(data["message"], str)64 65        # If successful, verify that files were processed66        if data["status"] == "completed":67            assert data["files_processed"] >= 1  # At least our test file68            assert data["chunks_created"] >= 1   # At least some chunks created69 70 71def test_ingestion_with_empty_directory(client, tmp_path):72    """Test ingestion with an empty directory."""73    # Create an empty temporary directory74    empty_dir = tmp_path / "empty_docs"75    empty_dir.mkdir()76 77    # Prepare ingestion request78    ingest_data = {79        "source_path": str(empty_dir)80    }81 82    # Call the ingestion endpoint83    response = client.post("/api/embeddings/ingest", json=ingest_data)84 85    # Should handle empty directory gracefully86    assert response.status_code == 20087    data = response.json()88 89    assert data["status"] in ["completed", "failed"]90    assert data["files_processed"] == 091    assert data["chunks_created"] == 092 93 94def test_ingestion_service_direct(client):95    """Test the ingestion service directly."""96    # Test the ingestion service with a mock path (it should handle non-existent paths gracefully)97    result = ingestion_service.ingest_book_content("nonexistent/path")98 99    # The service should return a structured response even for non-existent paths100    assert "status" in result101    assert "files_processed" in result102    assert "chunks_created" in result103    assert "message" in result104 105    assert result["status"] in ["completed", "failed"]106    assert isinstance(result["files_processed"], int)107    assert isinstance(result["chunks_created"], int)108    assert isinstance(result["message"], str)109 110 111def test_ingestion_pipeline_with_sample_content(client, tmp_path):112    """Test ingestion pipeline with sample book content."""113    # Create a more comprehensive test structure114    test_docs_dir = tmp_path / "docs"115    test_docs_dir.mkdir()116 117    # Create multiple test files118    file1 = test_docs_dir / "introduction.md"119    file1.write_text("# Introduction\nThis is the introduction to our book about AI and RAG systems.")120 121    file2 = test_docs_dir / "rag_basics.md"122    file2.write_text("# RAG Basics\nRetrieval-Augmented Generation combines retrieval and generation techniques.")123 124    file3 = test_docs_dir / "applications.md"125    file3.write_text("# Applications\nRAG is used in question answering, document understanding, and more.")126 127    # Prepare ingestion request128    ingest_data = {129        "source_path": str(test_docs_dir)130    }131 132    # Call the ingestion endpoint133    response = client.post("/api/embeddings/ingest", json=ingest_data)134 135    # Verify the response136    assert response.status_code == 200137    data = response.json()138 139    # Should have processed all 3 files140    assert data["status"] in ["completed", "failed"]141    assert isinstance(data["files_processed"], int)142    assert isinstance(data["chunks_created"], int)143    assert isinstance(data["message"], str)144 145    # If successful, should have processed at least 3 files146    if data["status"] == "completed":147        assert data["files_processed"] >= 3148        assert data["chunks_created"] >= 3  # At least one chunk per file