""" tests/test_smartrag.py Full test suite covering: - Unit tests : individual components (chunking, formatting, config) - Integration : RAG pipeline end-to-end with mock LLM - API tests : FastAPI endpoints via TestClient - Smoke test : quick sanity check without GPU Run all: pytest tests/ -v Run fast only: pytest tests/ -v -m "not slow" Run API only: pytest tests/test_smartrag.py::TestAPI -v """ import json import sys from pathlib import Path from unittest.mock import MagicMock, patch import pytest from fastapi.testclient import TestClient from langchain_core.documents import Document sys.path.insert(0, str(Path(__file__).parent.parent)) # ═══════════════════════════════════════════════════════════════════ # FIXTURES # ═══════════════════════════════════════════════════════════════════ @pytest.fixture(scope="session") def sample_docs(): """Sample documents for testing.""" return [ Document( page_content="Aspirin is a nonsteroidal anti-inflammatory drug (NSAID). " "It works by inhibiting COX-1 and COX-2 enzymes, reducing prostaglandin synthesis.", metadata={"source": "pharmacology_101.txt", "page": 1}, ), Document( page_content="Metformin is the first-line medication for type 2 diabetes. " "It works by decreasing hepatic glucose production and improving insulin sensitivity.", metadata={"source": "diabetes_guide.txt", "page": 1}, ), Document( page_content="The blood-brain barrier (BBB) is a selective semipermeable border " "of endothelial cells that prevents solutes in the circulating blood from " "non-selectively crossing into the extracellular fluid of the central nervous system.", metadata={"source": "neuroscience.txt", "page": 5}, ), ] @pytest.fixture def mock_rag_response(): """Mock RAGResponse for API testing.""" from rag.pipeline import RAGResponse return RAGResponse( question="What is aspirin?", answer="Aspirin is an NSAID that inhibits COX enzymes to reduce inflammation.", sources=["Source 1: pharmacology_101.txt"], context_used="Aspirin is a nonsteroidal anti-inflammatory drug...", num_chunks_retrieved=1, ) # ═══════════════════════════════════════════════════════════════════ # UNIT TESTS — Config # ═══════════════════════════════════════════════════════════════════ class TestConfig: def test_config_loads(self): from config import cfg assert cfg.model.base_model_id is not None assert cfg.rag.top_k > 0 assert cfg.rag.chunk_size > 0 def test_lora_config(self): from config import cfg assert cfg.lora.r > 0 assert cfg.lora.lora_alpha > 0 assert len(cfg.lora.target_modules) > 0 def test_training_config(self): from config import cfg assert 0 < cfg.training.learning_rate < 1 assert cfg.training.num_train_epochs > 0 def test_ensure_dirs_creates_directories(self, tmp_path, monkeypatch): from config import Config monkeypatch.setattr("config.cfg.model.output_dir", str(tmp_path / "model")) monkeypatch.setattr("config.cfg.rag.chroma_persist_dir", str(tmp_path / "chroma")) # Should not raise cfg_instance = Config() cfg_instance.ensure_dirs() # ═══════════════════════════════════════════════════════════════════ # UNIT TESTS — Data Preparation # ═══════════════════════════════════════════════════════════════════ class TestDataPreparation: def test_format_example_with_context(self): from data.prepare_dataset import format_example example = { "instruction": "What is aspirin?", "input": "Context about drugs", "output": "Aspirin is an NSAID.", } result = format_example(example) assert result is not None assert "[INST]" in result["text"] assert "[/INST]" in result["text"] assert "aspirin" in result["text"].lower() def test_format_example_without_context(self): from data.prepare_dataset import format_example example = { "instruction": "Explain photosynthesis", "input": "", "output": "Photosynthesis converts light to energy.", } result = format_example(example) assert result is not None assert "Context:" not in result["text"] def test_format_example_skips_empty(self): from data.prepare_dataset import format_example result = format_example({"instruction": "", "input": "", "output": ""}) assert result is None def test_clean_text(self): from data.prepare_dataset import clean_text dirty = " hello world \n\t " assert clean_text(dirty) == "hello world" # ═══════════════════════════════════════════════════════════════════ # UNIT TESTS — Vector Store (mocked embeddings) # ═══════════════════════════════════════════════════════════════════ class TestVectorStore: def test_chunk_documents(self, sample_docs): from rag.vectorstore import chunk_documents chunks = chunk_documents(sample_docs) assert len(chunks) >= len(sample_docs) for chunk in chunks: assert len(chunk.page_content) <= 600 # chunk_size + buffer def test_chunk_preserves_metadata(self, sample_docs): from rag.vectorstore import chunk_documents chunks = chunk_documents(sample_docs) # All chunks should have source metadata for chunk in chunks: assert "source" in chunk.metadata @patch("rag.vectorstore.HuggingFaceEmbeddings") @patch("rag.vectorstore.Chroma") def test_build_vectorstore(self, mock_chroma, mock_embeddings, sample_docs): from rag.vectorstore import build_vectorstore mock_chroma.from_documents.return_value = MagicMock() build_vectorstore(docs=sample_docs) mock_chroma.from_documents.assert_called_once() @patch("rag.vectorstore.HuggingFaceEmbeddings") @patch("rag.vectorstore.Chroma") def test_retrieve_returns_documents(self, mock_chroma, mock_embeddings, sample_docs): from rag.vectorstore import retrieve mock_store = MagicMock() mock_store.similarity_search_with_relevance_scores.return_value = [ (sample_docs[0], 0.9), (sample_docs[1], 0.7), ] results = retrieve("What is aspirin?", mock_store, top_k=2) assert len(results) == 2 assert results[0].page_content == sample_docs[0].page_content # ═══════════════════════════════════════════════════════════════════ # UNIT TESTS — RAG Pipeline (mocked LLM + vectorstore) # ═══════════════════════════════════════════════════════════════════ class TestRAGPipeline: @patch("rag.pipeline.load_vectorstore") @patch("rag.pipeline.load_finetuned_pipeline") def test_pipeline_query(self, mock_llm_loader, mock_vs_loader, sample_docs): from rag.pipeline import SmartRAGPipeline # Mock vectorstore mock_vs = MagicMock() mock_vs.similarity_search_with_relevance_scores.return_value = [ (sample_docs[0], 0.85) ] mock_vs_loader.return_value = mock_vs # Mock LLM pipeline mock_llm = MagicMock() mock_llm.return_value = [{"generated_text": "Aspirin inhibits COX enzymes."}] mock_llm_loader.return_value = mock_llm pipeline = SmartRAGPipeline() response = pipeline.query("What is aspirin?") assert response.question == "What is aspirin?" assert "aspirin" in response.answer.lower() or len(response.answer) > 0 assert response.num_chunks_retrieved >= 0 @patch("rag.pipeline.load_vectorstore") @patch("rag.pipeline.load_finetuned_pipeline") def test_pipeline_no_results(self, mock_llm_loader, mock_vs_loader): from rag.pipeline import SmartRAGPipeline mock_vs = MagicMock() mock_vs.similarity_search_with_relevance_scores.return_value = [] mock_vs_loader.return_value = mock_vs mock_llm_loader.return_value = MagicMock() pipeline = SmartRAGPipeline() response = pipeline.query("xyzzy nonsense query 12345") assert response.num_chunks_retrieved == 0 assert "couldn't find" in response.answer.lower() or len(response.answer) > 0 # ═══════════════════════════════════════════════════════════════════ # API TESTS — FastAPI endpoints # ═══════════════════════════════════════════════════════════════════ import api.app # ← ADD THIS LINE class TestAPI: @pytest.fixture def client(self, mock_rag_response): """Create test client with mocked pipeline.""" with patch("api.app.get_pipeline") as mock_get: mock_pipeline = MagicMock() mock_pipeline.query.return_value = mock_rag_response mock_pipeline.vectorstore = MagicMock() mock_get.return_value = mock_pipeline # Patch startup to avoid loading real models with patch("api.app.pipeline", mock_pipeline): from api.app import app yield TestClient(app) def test_health_endpoint(self, client): response = client.get("/health") assert response.status_code == 200 data = response.json() assert "status" in data assert "model_loaded" in data def test_root_endpoint(self, client): response = client.get("/") assert response.status_code == 200 assert "SmartRAG" in response.json()["name"] def test_query_endpoint_valid(self, client): response = client.post("/query", json={"question": "What is aspirin?"}) assert response.status_code == 200 data = response.json() assert "answer" in data assert "sources" in data assert "latency_ms" in data def test_query_endpoint_too_short(self, client): response = client.post("/query", json={"question": "hi"}) assert response.status_code == 422 # Pydantic validation error def test_query_endpoint_with_top_k(self, client): response = client.post("/query", json={"question": "What is aspirin?", "top_k": 3}) assert response.status_code == 200 def test_ingest_endpoint(self, client): response = client.post("/ingest", json={ "texts": ["Ibuprofen is an NSAID used for pain relief."], "metadata": [{"source": "test_doc.txt"}], }) assert response.status_code == 200 assert response.json()["status"] == "accepted" # ═══════════════════════════════════════════════════════════════════ # INTEGRATION TEST — End-to-end smoke test (no GPU needed) # ═══════════════════════════════════════════════════════════════════ @pytest.mark.slow class TestIntegration: """Integration tests that test more of the real stack (marked slow).""" def test_chunk_then_retrieve(self, sample_docs, tmp_path, monkeypatch): """Test chunking → embedding → retrieval pipeline (mocked embeddings).""" import numpy as np from rag.vectorstore import chunk_documents chunks = chunk_documents(sample_docs) assert len(chunks) > 0 # Verify chunk content integrity all_text = " ".join(c.page_content for c in chunks) assert "aspirin" in all_text.lower() assert "metformin" in all_text.lower() def test_data_pipeline_flow(self, tmp_path, monkeypatch): """Test data formatting → save → load round-trip.""" import json from data.prepare_dataset import format_example examples = [ {"instruction": "What is X?", "input": "Context X", "output": "X is great."}, {"instruction": "What is Y?", "input": "", "output": "Y is fine."}, ] formatted = [format_example(e) for e in examples] formatted = [f for f in formatted if f] # Save path = tmp_path / "test.jsonl" with open(path, "w") as f: for item in formatted: f.write(json.dumps(item) + "\n") # Load and verify loaded = [json.loads(line) for line in open(path)] assert len(loaded) == 2 assert all("text" in item for item in loaded) assert all("[INST]" in item["text"] for item in loaded)