tests/unit/test_triage_question_generator.py

#!/usr/bin/env python3
"""
Test script for TriageQuestionGenerator functionality.
"""

import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'src'))

from config.prompt_management.triage_question_generator import TriageQuestionGenerator
from config.prompt_management.data_models import ScenarioType, ConversationHistory

def test_triage_question_generator():
    """Test TriageQuestionGenerator functionality."""
    print("Testing TriageQuestionGenerator...")
    
    # Initialize generator
    generator = TriageQuestionGenerator()
    print("✓ TriageQuestionGenerator initialized")
    
    # Test 1: Scenario identification
    print("\n1. Testing scenario identification...")
    
    test_statements = [
        ("I used to love gardening, but now I can't", ScenarioType.LOSS_OF_INTEREST),
        ("My mother passed away last month", ScenarioType.LOSS_OF_LOVED_ONE),
        ("I don't have anyone to help me", ScenarioType.NO_SUPPORT),
        ("I feel some stress", ScenarioType.VAGUE_STRESS),
        ("I can't sleep at night", ScenarioType.SLEEP_ISSUES)
    ]
    
    for statement, expected_type in test_statements:
        identified_type = generator.identify_scenario_type(statement)
        assert identified_type == expected_type, f"'{statement}' → Expected {expected_type.value}, got {identified_type}"
        print(f"   ✓ '{statement}' → {expected_type.value}")
    
    # Test 2: Scenario creation from statements
    print("\n2. Testing scenario creation...")
    
    for statement, expected_type in test_statements:
        scenario = generator.create_scenario_from_statement(statement)
        assert scenario is not None, f"Failed to create scenario for: {statement}"
        assert scenario.scenario_type == expected_type, f"Wrong scenario type for: {statement}"
        print(f"   ✓ Created scenario for: {expected_type.value}")
        print(f"     Context clues: {len(scenario.context_clues)}")
        print(f"     Question patterns: {len(scenario.question_patterns)}")
    
    # Test 3: Targeted question generation
    print("\n3. Testing targeted question generation...")
    
    for statement, expected_type in test_statements:
        scenario = generator.create_scenario_from_statement(statement)
        assert scenario is not None, f"No scenario created for: {statement}"
        
        question = generator.generate_targeted_question(scenario)
        print(f"   ✓ {expected_type.value}:")
        print(f"     Statement: {statement}")
        print(f"     Question: {question}")
        
        # Validate question is not empty and is a question
        assert len(question.strip()) > 0, "Empty question generated"
        assert question.strip().endswith('?'), "Generated text is not a question"
    
    # Test 4: Question effectiveness validation
    print("\n4. Testing question effectiveness validation...")
    
    test_questions = [
        ("Is that something that's been weighing on you emotionally, or is it more about circumstances?", "loss_of_interest", 0.7),
        ("How are you feeling?", "loss_of_interest", 0.3),
        ("I'm sorry for your loss. How have you been coping with this?", "loss_of_loved_one", 0.7),
        ("That's sad.", "loss_of_loved_one", 0.2)
    ]
    
    for question, scenario, min_expected_score in test_questions:
        score = generator.validate_question_effectiveness(question, scenario)
        if score >= min_expected_score:
            print(f"   ✓ '{question[:40]}...' → Score: {score:.2f}")
        else:
            print(f"   ⚠ '{question[:40]}...' → Score: {score:.2f} (expected >= {min_expected_score})")
    
    # Test 5: Question patterns retrieval
    print("\n5. Testing question patterns retrieval...")
    
    for scenario_type in ScenarioType:
        patterns = generator.get_question_patterns(scenario_type.value)
        print(f"   ✓ {scenario_type.value}: {len(patterns)} patterns")
        
        if patterns:
            sample_pattern = patterns[0]
            print(f"     Sample: {sample_pattern.template[:60]}...")
    
    # Test 6: Variable extraction and template rendering
    print("\n6. Testing variable extraction...")
    
    test_cases = [
        ("I used to love gardening, but now I can't", ScenarioType.LOSS_OF_INTEREST),
        ("My mother passed away", ScenarioType.LOSS_OF_LOVED_ONE),
        ("I feel stressed", ScenarioType.VAGUE_STRESS)
    ]
    
    for statement, scenario_type in test_cases:
        patterns = generator._scenario_patterns.get(scenario_type, [])
        if patterns:
            variables = generator._extract_variables(statement, patterns[0])
            print(f"   ✓ '{statement}' → Variables: {variables}")
        else:
            print(f"   ⚠ No patterns for {scenario_type.value}")
    
    # Test 7: Fallback question generation
    print("\n7. Testing fallback question generation...")
    
    fallback_statements = [
        "Something is wrong",
        "I don't know what to do",
        "This is confusing"
    ]
    
    for statement in fallback_statements:
        fallback_question = generator._generate_fallback_question(statement)
        print(f"   ✓ '{statement}' → '{fallback_question}'")
        assert fallback_question.endswith('?'), "Fallback is not a question"
    
    # Test 8: Context integration
    print("\n8. Testing context integration...")
    
    # Create mock conversation history
    from config.prompt_management.data_models import Message
    from datetime import datetime
    context = ConversationHistory(
        messages=[Message(content="Previous message", classification="yellow", timestamp=datetime.fromisoformat("2024-01-01T00:00:00"))],
        distress_indicators_found=["sleep_difficulties"],
        context_flags=["medical_context", "previous_distress"]
    )
    
    statement = "I can't sleep"
    scenario = generator.create_scenario_from_statement(statement, context)
    
    assert scenario is not None, "Failed to integrate context"
    print(f"   ✓ Context integrated: {len(scenario.context_clues)} clues")
    print(f"     Context clues: {scenario.context_clues}")
    
    print("\n✓ All TriageQuestionGenerator tests passed!")

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