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!")