run_tests.py

import os
import sys
import json

# Ensure parent directory is in python path
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

from backend.services import csv_service, sql_service, llm_service, insight_service

def test_csv_ingestion():
    print("\n=== Test 1: CSV Ingestion & Sample Datasets ===")
    try:
        # get_sample_datasets automatically loads files if not loaded
        datasets = csv_service.get_sample_datasets()
        print(f"Discovered and loaded {len(datasets)} datasets.")
        for ds in datasets:
            print(f"- Name: {ds['name']}, Rows: {ds['row_count']}, Columns count: {len(ds['columns'])}")
            
        # Verify specific datasets exist in the list
        names = [d["name"] for d in datasets]
        assert "sales" in names, "sales dataset missing"
        assert "employees" in names, "employees dataset missing"
        assert "ecommerce" in names, "ecommerce dataset missing"
        
        print("Test 1 passed successfully!")
        return datasets
    except Exception as e:
        print(f"Test 1 failed: {e}")
        sys.exit(1)

def test_sql_validation():
    print("\n=== Test 2: SQL Safety Validation ===")
    try:
        # Safe query
        safe_sql = "SELECT employee_id, name, salary FROM data WHERE department = 'Engineering' ORDER BY salary DESC;"
        assert sql_service.validate_sql(safe_sql) == True, "Safe query flagged as dangerous"
        print("[OK] Safe SELECT validated successfully")

        # Dangerous queries
        unsafe_queries = [
            "DROP TABLE data;",
            "INSERT INTO data (name) VALUES ('Hacker');",
            "UPDATE data SET salary = 999999;",
            "DELETE FROM data;",
            "CREATE TABLE hack (id int);",
            "ALTER TABLE data ADD COLUMN hack TEXT;",
            "EXEC xp_cmdshell 'whoami';"
        ]
        
        for q in unsafe_queries:
            assert sql_service.validate_sql(q) == False, f"Dangerous query allowed: {q}"
            print(f"[OK] Blocked dangerous query: {q}")

        print("Test 2 passed successfully!")
    except Exception as e:
        print(f"Test 2 failed: {e}")
        sys.exit(1)

def test_end_to_end_pipeline(datasets):
    print("\n=== Test 3: End-to-End NL to SQL to Insight ===")
    try:
        # Find employees dataset schema
        emp_ds = next(d for d in datasets if d["name"] == "employees")
        schema = {
            "table_name": "data",
            "columns": emp_ds["columns"],
            "row_count": emp_ds["row_count"]
        }
        
        question = "What is the average salary and performance score in the Engineering department?"
        print(f"Question: '{question}'")
        
        # 1. Generate SQL
        gen_result = llm_service.generate_sql(question, schema)
        generated_sql = gen_result["sql"]
        print(f"Generated SQL: {generated_sql}")
        
        # 2. Execute SQL
        exec_result = sql_service.execute_query("employees", generated_sql)
        print(f"Execution Results:\nColumns: {exec_result['columns']}")
        print(f"Rows: {exec_result['rows']}")
        print(f"Time: {exec_result['execution_time_ms']} ms")
        
        # 3. Generate Insight
        insight = insight_service.generate_insight(question, generated_sql, exec_result)
        print(f"Generated Insight:\n{insight}")
        
        print("Test 3 passed successfully!")
    except Exception as e:
        print(f"Test 3 failed: {e}")
        sys.exit(1)

def test_sql_fixing(datasets):
    print("\n=== Test 4: SQL Auto-Fixing on Execution Error ===")
    try:
        emp_ds = next(d for d in datasets if d["name"] == "employees")
        schema = {
            "table_name": "data",
            "columns": emp_ds["columns"],
            "row_count": emp_ds["row_count"]
        }
        
        # A bad query that references a column 'departmet' (typo) instead of 'department'
        bad_sql = "SELECT name, salary FROM data WHERE departmet = 'Engineering' LIMIT 5;"
        print(f"Attempting execution of bad SQL: {bad_sql}")
        
        try:
            sql_service.execute_query("employees", bad_sql)
            print("Error: Bad SQL was expected to fail execution but succeeded.")
            sys.exit(1)
        except ValueError as err:
            error_message = str(err)
            print(f"Caught expected SQLite error: '{error_message}'")
            
            # Now ask the LLM to fix it
            fix_result = llm_service.fix_sql(bad_sql, error_message, schema)
            fixed_sql = fix_result["sql"]
            print(f"Fixed SQL generated by LLM: {fixed_sql}")
            
            # Run the fixed SQL
            exec_result = sql_service.execute_query("employees", fixed_sql)
            print(f"Execution of Fixed SQL succeeded! Rows returned: {exec_result['row_count']}")
            print(f"First row: {exec_result['rows'][0] if exec_result['rows'] else 'None'}")
            
        print("Test 4 passed successfully!")
    except Exception as e:
        print(f"Test 4 failed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    print("Starting tests for Natural Language Data Analyst backend...")
    datasets = test_csv_ingestion()
    test_sql_validation()
    test_end_to_end_pipeline(datasets)
    test_sql_fixing(datasets)
    print("\nAll tests passed successfully!")