backend/test_prompt_templates.py

"""
Unit Tests for Medical Prompt Templates
Tests prompt generation logic without requiring ML model dependencies

Author: MiniMax Agent
Date: 2025-10-29
"""

import sys
sys.path.insert(0, '/workspace/medical-ai-platform/backend')

from medical_prompt_templates import PromptTemplateLibrary, SummaryType


def create_sample_ecg_data():
    """Sample ECG data for testing"""
    return {
        "metadata": {
            "document_id": "ecg-test-001",
            "facility": "Test Hospital",
            "document_date": "2025-10-29"
        },
        "intervals": {
            "pr_ms": 165.0,
            "qrs_ms": 92.0,
            "qt_ms": 390.0,
            "qtc_ms": 425.0,
            "rr_ms": 850.0
        },
        "rhythm_classification": {
            "primary_rhythm": "Normal Sinus Rhythm",
            "heart_rate_bpm": 71,
            "heart_rate_regularity": "regular",
            "arrhythmia_types": []
        },
        "arrhythmia_probabilities": {
            "normal_rhythm": 0.92,
            "atrial_fibrillation": 0.02
        },
        "derived_features": {
            "st_elevation_mm": {},
            "axis_deviation": "normal",
            "t_wave_abnormalities": []
        }
    }


def create_sample_model_outputs():
    """Sample model outputs"""
    return [
        {
            "model_name": "Bio_ClinicalBERT",
            "domain": "clinical_notes",
            "result": {"summary": "Analysis complete", "confidence": 0.87}
        }
    ]


def test_ecg_clinician_prompt():
    """Test ECG clinician prompt generation"""
    print("\n" + "="*80)
    print("TEST 1: ECG Clinician Prompt Generation")
    print("="*80)
    
    lib = PromptTemplateLibrary()
    ecg_data = create_sample_ecg_data()
    model_outputs = create_sample_model_outputs()
    confidence = {"overall_confidence": 0.89}
    
    prompt = lib.get_clinician_summary_template(
        modality="ECG",
        structured_data=ecg_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    # Validate prompt contains key elements
    assert "ECG" in prompt, "Prompt should mention ECG"
    assert "Heart Rate: 71 bpm" in prompt, "Prompt should include heart rate"
    assert "Normal Sinus Rhythm" in prompt, "Prompt should include rhythm"
    assert "ANALYSIS CONFIDENCE: 89.0%" in prompt, "Prompt should include confidence"
    assert "TECHNICAL SUMMARY" in prompt, "Prompt should have technical summary section"
    assert "RECOMMENDATIONS" in prompt, "Prompt should have recommendations section"
    
    print(f"✓ Prompt generated: {len(prompt)} characters")
    print(f"✓ Contains all required sections")
    print(f"\nSample excerpt:\n{prompt[:300]}...\n")
    
    return True


def test_ecg_patient_prompt():
    """Test ECG patient-friendly prompt generation"""
    print("\n" + "="*80)
    print("TEST 2: ECG Patient Prompt Generation")
    print("="*80)
    
    lib = PromptTemplateLibrary()
    ecg_data = create_sample_ecg_data()
    model_outputs = create_sample_model_outputs()
    confidence = {"overall_confidence": 0.89}
    
    prompt = lib.get_patient_summary_template(
        modality="ECG",
        structured_data=ecg_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    # Validate patient-friendly language
    assert "YOUR ECG RESULTS" in prompt, "Should use patient-friendly heading"
    assert "simple" in prompt.lower(), "Should mention simple language"
    assert "WHAT WE FOUND" in prompt, "Should have clear sections"
    assert "NEXT STEPS" in prompt, "Should include next steps"
    
    print(f"✓ Prompt generated: {len(prompt)} characters")
    print(f"✓ Patient-friendly language detected")
    print(f"\nSample excerpt:\n{prompt[:300]}...\n")
    
    return True


def test_radiology_prompts():
    """Test radiology prompt generation"""
    print("\n" + "="*80)
    print("TEST 3: Radiology Prompt Generation")
    print("="*80)
    
    lib = PromptTemplateLibrary()
    rad_data = {
        "metadata": {"document_id": "rad-001"},
        "image_references": [
            {"modality": "CT", "body_part": "Chest"}
        ],
        "findings": {
            "findings_text": "Clear lungs bilaterally",
            "impression_text": "No acute abnormality",
            "critical_findings": [],
            "incidental_findings": []
        },
        "metrics": {"organ_volumes": {}, "lesion_measurements": []}
    }
    model_outputs = create_sample_model_outputs()
    confidence = {"overall_confidence": 0.85}
    
    # Test clinician prompt
    clinician_prompt = lib.get_clinician_summary_template(
        modality="radiology",
        structured_data=rad_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    assert "IMAGING STUDY DETAILS" in clinician_prompt
    assert "CT" in clinician_prompt
    assert "Chest" in clinician_prompt
    
    # Test patient prompt
    patient_prompt = lib.get_patient_summary_template(
        modality="radiology",
        structured_data=rad_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    assert "YOUR IMAGING STUDY" in patient_prompt
    assert "Type of Scan" in patient_prompt
    
    print(f"✓ Clinician prompt: {len(clinician_prompt)} characters")
    print(f"✓ Patient prompt: {len(patient_prompt)} characters")
    
    return True


def test_laboratory_prompts():
    """Test laboratory prompt generation"""
    print("\n" + "="*80)
    print("TEST 4: Laboratory Prompt Generation")
    print("="*80)
    
    lib = PromptTemplateLibrary()
    lab_data = {
        "metadata": {"document_id": "lab-001"},
        "tests": [
            {
                "test_name": "Glucose",
                "value": 105.0,
                "unit": "mg/dL",
                "reference_range_low": 70.0,
                "reference_range_high": 99.0,
                "flags": ["H"]
            }
        ],
        "abnormal_count": 1,
        "critical_values": [],
        "panel_name": "Basic Metabolic Panel",
        "collection_date": "2025-10-29"
    }
    model_outputs = create_sample_model_outputs()
    confidence = {"overall_confidence": 0.92}
    
    # Test clinician prompt
    clinician_prompt = lib.get_clinician_summary_template(
        modality="laboratory",
        structured_data=lab_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    assert "LABORATORY PANEL" in clinician_prompt
    assert "Glucose" in clinician_prompt
    assert "105.0" in clinician_prompt
    assert "SUMMARY OF KEY FINDINGS" in clinician_prompt
    
    # Test patient prompt
    patient_prompt = lib.get_patient_summary_template(
        modality="laboratory",
        structured_data=lab_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    assert "YOUR LAB RESULTS" in patient_prompt
    assert "everyday language" in patient_prompt.lower()
    
    print(f"✓ Clinician prompt: {len(clinician_prompt)} characters")
    print(f"✓ Patient prompt: {len(patient_prompt)} characters")
    
    return True


def test_clinical_notes_prompts():
    """Test clinical notes prompt generation"""
    print("\n" + "="*80)
    print("TEST 5: Clinical Notes Prompt Generation")
    print("="*80)
    
    lib = PromptTemplateLibrary()
    notes_data = {
        "metadata": {"document_id": "note-001"},
        "note_type": "progress_note",
        "sections": [
            {
                "section_type": "chief_complaint",
                "content": "Patient presents with chest pain"
            }
        ],
        "entities": [],
        "diagnoses": ["Chest pain, unspecified"],
        "medications": ["Aspirin 81mg daily"]
    }
    model_outputs = create_sample_model_outputs()
    confidence = {"overall_confidence": 0.87}
    
    # Test clinician prompt
    clinician_prompt = lib.get_clinician_summary_template(
        modality="clinical_notes",
        structured_data=notes_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    assert "CLINICAL SECTIONS" in clinician_prompt
    assert "ASSESSMENT" in clinician_prompt
    assert "chest pain" in clinician_prompt.lower()
    
    # Test patient prompt
    patient_prompt = lib.get_patient_summary_template(
        modality="clinical_notes",
        structured_data=notes_data,
        model_outputs=model_outputs,
        confidence_scores=confidence
    )
    
    assert "REASON FOR YOUR VISIT" in patient_prompt
    assert "TREATMENT PLAN" in patient_prompt
    
    print(f"✓ Clinician prompt: {len(clinician_prompt)} characters")
    print(f"✓ Patient prompt: {len(patient_prompt)} characters")
    
    return True


def test_multi_modal_prompt():
    """Test multi-modal synthesis prompt"""
    print("\n" + "="*80)
    print("TEST 6: Multi-Modal Synthesis Prompt")
    print("="*80)
    
    lib = PromptTemplateLibrary()
    
    modalities = ["ECG", "radiology", "laboratory"]
    all_data = {
        "ECG": create_sample_ecg_data(),
        "radiology": {"metadata": {"document_id": "rad-001"}},
        "laboratory": {"metadata": {"document_id": "lab-001"}}
    }
    confidence_scores = {
        "ECG": 0.89,
        "radiology": 0.85,
        "laboratory": 0.92
    }
    
    prompt = lib.get_multi_modal_synthesis_template(
        modalities=modalities,
        all_data=all_data,
        confidence_scores=confidence_scores
    )
    
    assert "multiple medical documents" in prompt.lower()
    assert "ECG" in prompt
    assert "INTEGRATED CLINICAL PICTURE" in prompt
    assert "COORDINATED CARE PLAN" in prompt
    
    print(f"✓ Multi-modal prompt: {len(prompt)} characters")
    print(f"✓ Includes all {len(modalities)} modalities")
    
    return True


def test_confidence_explanation_prompt():
    """Test confidence explanation prompt"""
    print("\n" + "="*80)
    print("TEST 7: Confidence Explanation Prompt")
    print("="*80)
    
    lib = PromptTemplateLibrary()
    
    # Test high confidence
    high_conf = {
        "overall_confidence": 0.92,
        "extraction_confidence": 0.94,
        "model_confidence": 0.91,
        "data_quality": 0.95
    }
    
    prompt_high = lib.get_confidence_explanation_template(
        confidence_scores=high_conf,
        modality="ECG"
    )
    
    assert "92.0%" in prompt_high
    assert "AUTO-APPROVED" in prompt_high
    
    # Test low confidence
    low_conf = {
        "overall_confidence": 0.55,
        "extraction_confidence": 0.55,
        "model_confidence": 0.50,
        "data_quality": 0.58
    }
    
    prompt_low = lib.get_confidence_explanation_template(
        confidence_scores=low_conf,
        modality="ECG"
    )
    
    assert "55.0%" in prompt_low
    assert "MANUAL REVIEW REQUIRED" in prompt_low
    
    print(f"✓ High confidence prompt generated")
    print(f"✓ Low confidence prompt generated")
    print(f"✓ Threshold detection working correctly")
    
    return True


def run_prompt_template_tests():
    """Run all prompt template tests"""
    print("\n" + "="*80)
    print("MEDICAL PROMPT TEMPLATES - UNIT TEST SUITE")
    print("Testing Prompt Generation Logic")
    print("="*80)
    
    tests = [
        ("ECG Clinician Prompt", test_ecg_clinician_prompt),
        ("ECG Patient Prompt", test_ecg_patient_prompt),
        ("Radiology Prompts", test_radiology_prompts),
        ("Laboratory Prompts", test_laboratory_prompts),
        ("Clinical Notes Prompts", test_clinical_notes_prompts),
        ("Multi-Modal Prompt", test_multi_modal_prompt),
        ("Confidence Explanation", test_confidence_explanation_prompt)
    ]
    
    results = []
    
    for test_name, test_func in tests:
        try:
            success = test_func()
            results.append((test_name, "PASS" if success else "FAIL"))
            print(f"✓ {test_name}: PASS")
        except AssertionError as e:
            print(f"✗ {test_name}: FAIL - {str(e)}")
            results.append((test_name, "FAIL"))
        except Exception as e:
            print(f"✗ {test_name}: ERROR - {str(e)}")
            import traceback
            traceback.print_exc()
            results.append((test_name, "ERROR"))
    
    # Print summary
    print("\n" + "="*80)
    print("TEST SUMMARY")
    print("="*80)
    for test_name, status in results:
        status_symbol = "✓" if status == "PASS" else "✗"
        print(f"{status_symbol} {test_name}: {status}")
    
    passed = sum(1 for _, status in results if status == "PASS")
    total = len(results)
    print(f"\nTotal: {passed}/{total} tests passed ({passed/total*100:.1f}%)")
    print("="*80)
    
    return passed == total


if __name__ == "__main__":
    success = run_prompt_template_tests()
    exit(0 if success else 1)