test_production_engine.py

#!/usr/bin/env python3
"""
COMPREHENSIVE TEST SUITE for Production-Grade NLP Engine
Tests all query types and filter combinations
Run: python test_production_engine.py
"""

import sys
from database import FilterExtractor, ask_llm

def test_basic_filters():
    """Test basic single-filter extraction"""
    print("\n" + "="*70)
    print("TEST 1: BASIC FILTERS")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        ("show TN vehicles", {"state": "TN"}),
        ("show buses", {"vehicle_type": "bus"}),
        ("show vehicles in adyar", {"location": "adyar"}),
        ("track TN63MB3157", {"plate": "TN63MB3157"}),
    ]
    
    for query, expected_filters in test_cases:
        filters = extractor.extract_filters(query)
        status = "✅" if all(
            filters.get(k) == v for k, v in expected_filters.items()
        ) else "❌"
        print(f"{status} {query}")
        print(f"   Extracted: {expected_filters}")


def test_multi_filters():
    """Test multi-filter extraction"""
    print("\n" + "="*70)
    print("TEST 2: MULTI-FILTER COMBINATIONS")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        ("show TN buses in adyar", 
         {"state": "TN", "vehicle_type": "bus", "location": "adyar"}),
        ("show Karnataka cars in velachery",
         {"state": "KA", "vehicle_type": "car", "location": "velachery"}),
        ("show TN buses in besant nagar on monday",
         {"state": "TN", "vehicle_type": "bus", "location": "besant nagar", "day": "Monday"}),
    ]
    
    for query, expected_filters in test_cases:
        filters = extractor.extract_filters(query)
        status = "✅" if all(
            filters.get(k) == v for k, v in expected_filters.items()
        ) else "❌"
        print(f"{status} {query}")


def test_date_range():
    """Test date range extraction"""
    print("\n" + "="*70)
    print("TEST 3: DATE RANGE EXTRACTION")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        "show vehicles from 01-05-2026 to 10-05-2026",
        "show buses between 2026-05-01 and 2026-05-10",
        "show vehicles from 01/05/2026 to 10/05/2026",
    ]
    
    for query in test_cases:
        filters = extractor.extract_filters(query)
        date_range = filters.get("date_range")
        if date_range:
            print(f"✅ {query}")
            print(f"   Start: {date_range['start']}, End: {date_range['end']}")
        else:
            print(f"❌ {query}")


def test_time_range():
    """Test time range extraction"""
    print("\n" + "="*70)
    print("TEST 4: TIME RANGE EXTRACTION")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        ("show vehicles after 8 PM", {"start": 20, "end": 23}),
        ("show vehicles before 6 AM", {"start": 0, "end": 6}),
        ("show vehicles between 6 PM and 9 PM", {"start": 18, "end": 21}),
        ("show vehicles in the morning", {"start": 5, "end": 12}),
        ("show vehicles in the evening", {"start": 17, "end": 21}),
    ]
    
    for query, expected_range in test_cases:
        filters = extractor.extract_filters(query)
        time_range = filters.get("time_range")
        if time_range:
            if (time_range.get("start") == expected_range["start"] and
                time_range.get("end") == expected_range["end"]):
                print(f"✅ {query}")
            else:
                print(f"⚠️ {query}")
                print(f"   Expected: {expected_range}, Got: {time_range}")
        else:
            print(f"❌ {query}")


def test_intent_detection():
    """Test intent detection"""
    print("\n" + "="*70)
    print("TEST 5: INTENT DETECTION")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        ("track TN63MB3157", ["tracking"]),
        ("count buses", ["count"]),
        ("show top vehicles", ["top", "analytics"]),
        ("show hourly traffic", ["hourly", "analytics"]),
        ("show suspicious vehicles", ["suspicious", "analytics"]),
        ("show vehicles in multiple locations", ["location_based"]),
    ]
    
    for query, expected_intents in test_cases:
        intents = extractor.detect_intents(query)
        detected = [k for k, v in intents.items() if v]
        
        match = all(intent in detected for intent in expected_intents)
        status = "✅" if match else "⚠️"
        
        print(f"{status} {query}")
        print(f"   Intents: {', '.join(detected) if detected else 'general'}")


def test_sql_generation():
    """Test SQL generation for various queries"""
    print("\n" + "="*70)
    print("TEST 6: SQL GENERATION")
    print("="*70)
    
    test_cases = [
        "show TN buses",
        "show buses in adyar",
        "show TN buses in adyar from 01-05-2026 to 10-05-2026",
        "show vehicles after 8 PM",
        "show buses between 6 PM and 9 PM",
        "count TN vehicles",
        "track TN63MB3157",
        "show top vehicles",
        "show hourly traffic",
        "show suspicious vehicles",
    ]
    
    for query in test_cases:
        sql = ask_llm(query)
        # Validate SQL
        if sql.strip().startswith("SELECT"):
            print(f"✅ {query}")
            # Show first line of SQL
            first_line = sql.split("\n")[0]
            print(f"   {first_line}...")
        else:
            print(f"❌ {query}")
            print(f"   Invalid SQL: {sql[:50]}...")


def test_complex_queries():
    """Test complex multi-dimension queries"""
    print("\n" + "="*70)
    print("TEST 7: COMPLEX MULTI-DIMENSION QUERIES")
    print("="*70)
    
    complex_queries = [
        "show TN buses in adyar from 01-05-2026 to 10-05-2026 after 8 PM",
        "show TN high-confidence buses detected in multiple locations",
        "show Kerala cars in velachery on weekend between 6 PM and 9 PM",
        "count vehicles in adyar from 01-05-2026 to 10-05-2026",
        "track TN63MB3157 in adyar on monday",
    ]
    
    for query in complex_queries:
        sql = ask_llm(query)
        if "SELECT" in sql and "FROM vehicle_logs" in sql:
            print(f"✅ {query}")
        else:
            print(f"❌ {query}")


def test_location_variants():
    """Test location variant matching"""
    print("\n" + "="*70)
    print("TEST 8: LOCATION VARIANT MATCHING")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        ("adyar", "adyar"),
        ("besant nagar", "besant nagar"),
        ("besant", "besant nagar"),
        ("t nagar", "t nagar"),
        ("tnagar", "t nagar"),
        ("anna nagar", "anna nagar"),
        ("anna", "anna nagar"),
        ("velachery", "velachery"),
    ]
    
    for query_location, expected in test_cases:
        filters = extractor.extract_filters(f"show vehicles in {query_location}")
        location = filters.get("location")
        status = "✅" if location == expected else "❌"
        print(f"{status} '{query_location}' → '{location}'")


def test_vehicle_synonyms():
    """Test vehicle type synonym matching"""
    print("\n" + "="*70)
    print("TEST 9: VEHICLE TYPE SYNONYMS")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        ("buses", "bus"),
        ("truck", "truck"),
        ("lorry", "truck"),
        ("motorcycle", "bike"),
        ("motorbike", "bike"),
        ("autorickshaw", "auto"),
        ("auto-rickshaw", "auto"),
        ("compact", "car"),
        ("sedan", "car"),
        ("cabs", "taxi"),
    ]
    
    for query_type, expected in test_cases:
        filters = extractor.extract_filters(f"show {query_type}")
        vehicle_type = filters.get("vehicle_type")
        status = "✅" if vehicle_type == expected else "❌"
        print(f"{status} '{query_type}' → '{vehicle_type}'")


def test_confidence_threshold():
    """Test confidence threshold extraction"""
    print("\n" + "="*70)
    print("TEST 10: CONFIDENCE THRESHOLD")
    print("="*70)
    
    extractor = FilterExtractor()
    
    test_cases = [
        ("show vehicles with 0.9 confidence", 0.9),
        ("show high confidence detections above 0.95", 0.95),
        ("vehicles with 0.85 confidence", 0.85),
    ]
    
    for query, expected_conf in test_cases:
        filters = extractor.extract_filters(query)
        confidence = filters.get("confidence")
        if confidence:
            status = "✅" if abs(confidence - expected_conf) < 0.01 else "⚠️"
            print(f"{status} {query}")
            print(f"   Confidence: {confidence}")
        else:
            print(f"⚠️ {query}")
            print(f"   Confidence not extracted")


def run_all_tests():
    """Run all test suites"""
    print("\n\n")
    print("╔" + "═"*68 + "╗")
    print("║" + " "*68 + "║")
    print("║" + "  PRODUCTION-GRADE NLP ENGINE TEST SUITE  ".center(68) + "║")
    print("║" + " "*68 + "║")
    print("╚" + "═"*68 + "╝")
    
    try:
        test_basic_filters()
        test_multi_filters()
        test_date_range()
        test_time_range()
        test_intent_detection()
        test_sql_generation()
        test_complex_queries()
        test_location_variants()
        test_vehicle_synonyms()
        test_confidence_threshold()
        
        print("\n" + "="*70)
        print("✅ ALL TESTS COMPLETED")
        print("="*70 + "\n")
        
    except Exception as e:
        print(f"\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    run_all_tests()