ruby2210/rag-chatbot
0
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