simple_test.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Simple SafeRAG Test
Basic functionality test without complex dependencies
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

def test_imports():
    """Test that all modules can be imported"""
    print("Testing imports...")
    
    try:
        from data_processing import DataLoader, Preprocessor
        print("+ DataLoader and Preprocessor imported successfully")
    except Exception as e:
        print("✗ Failed to import DataLoader/Preprocessor:", e)
        return False
    
    try:
        from retriever import Embedder, FAISSIndex, Retriever, Reranker
        print("+ Retriever modules imported successfully")
    except Exception as e:
        print("✗ Failed to import retriever modules:", e)
        return False
    
    try:
        from generator import VLLMServer, SafeGenerator, PromptTemplates
        print("+ Generator modules imported successfully")
    except Exception as e:
        print("✗ Failed to import generator modules:", e)
        return False
    
    try:
        from calibration import RiskFeatureExtractor, CalibrationHead
        print("+ Calibration modules imported successfully")
    except Exception as e:
        print("✗ Failed to import calibration modules:", e)
        return False
    
    try:
        from eval import QAEvaluator, AttributionEvaluator, CalibrationEvaluator
        print("+ Evaluation modules imported successfully")
    except Exception as e:
        print("✗ Failed to import evaluation modules:", e)
        return False
    
    return True

def test_basic_functionality():
    """Test basic functionality without heavy dependencies"""
    print("\nTesting basic functionality...")
    
    try:
        # Test Preprocessor
        from data_processing.preprocessor import Preprocessor
        preprocessor = Preprocessor()
        
        # Test text cleaning
        text = "  This is a test text.   "
        cleaned = preprocessor.clean_text(text)
        assert cleaned == "This is a test text.", "Expected 'This is a test text.', got '{}'".format(cleaned)
        print("+ Text cleaning works")
        
        # Test sentence extraction
        text = "First sentence. Second sentence. Third sentence."
        sentences = preprocessor.extract_sentences(text)
        assert len(sentences) == 3, "Expected 3 sentences, got {}".format(len(sentences))
        print("+ Sentence extraction works")
        
    except Exception as e:
        print("✗ Preprocessor test failed:", e)
        return False
    
    try:
        # Test PromptTemplates
        from generator.prompt_templates import PromptTemplates
        templates = PromptTemplates()
        
        # Test prompt formatting
        prompt = templates.format_prompt(
            'rag',
            question="What is AI?",
            context="AI is artificial intelligence."
        )
        assert "What is AI?" in prompt, "Question not found in prompt"
        assert "AI is artificial intelligence." in prompt, "Context not found in prompt"
        print("+ Prompt templates work")
        
    except Exception as e:
        print("✗ PromptTemplates test failed:", e)
        return False
    
    try:
        # Test QAEvaluator
        from eval.eval_qa import QAEvaluator
        evaluator = QAEvaluator()
        
        # Test exact match
        predictions = ["Paris", "Paris"]
        references = ["Paris", "London"]
        em = evaluator.exact_match(predictions, references)
        assert em == 0.5, "Expected 0.5, got {}".format(em)
        print("+ QA evaluation works")
        
    except Exception as e:
        print("✗ QAEvaluator test failed:", e)
        return False
    
    return True

def test_config():
    """Test configuration loading"""
    print("\nTesting configuration...")
    
    try:
        import yaml
        with open('config.yaml', 'r') as f:
            config = yaml.safe_load(f)
        
        # Check required sections
        required_sections = ['models', 'data', 'index', 'retrieval', 'calibration', 'evaluation']
        for section in required_sections:
            assert section in config, "Missing config section: {}".format(section)
        
        print("+ Configuration file is valid")
        return True
        
    except Exception as e:
        print("✗ Configuration test failed:", e)
        return False

def main():
    """Run all tests"""
    print("SafeRAG Simple Test Suite")
    print("=" * 40)
    
    all_passed = True
    
    # Test imports
    if not test_imports():
        all_passed = False
    
    # Test basic functionality
    if not test_basic_functionality():
        all_passed = False
    
    # Test configuration
    if not test_config():
        all_passed = False
    
    print("\n" + "=" * 40)
    if all_passed:
        print("+ All tests passed!")
        print("SafeRAG is ready to use.")
    else:
        print("✗ Some tests failed.")
        print("Please check the errors above.")
    
    return all_passed

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)