Spaces:
Runtime error
Runtime error
File size: 5,435 Bytes
37a6ee1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | 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!")
|