test_smoke.py

"""
End-to-end smoke test for the Zero-LLM engine.
Creates a sample CSV, uploads it, runs commands, verifies results.
"""
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import polars as pl
from services.file_manager import handle_upload, delete_session
from services.session_manager import session_manager
from core.column_registry import column_registry
from core.intent_parser import parse_intent
from core.validator import validate_intent
from core.router import dispatch


def create_test_csv():
    df = pl.DataFrame({
        "Name": ["Rahul", "Priya", "Amit", "Sneha", "Vikram", "Neha", "Ravi", "Pooja"],
        "City": ["Mumbai", "Delhi", "Mumbai", "Bangalore", "Delhi", "Mumbai", "Bangalore", "Delhi"],
        "Salary": [50000, 60000, 55000, 70000, 65000, 48000, 72000, 58000],
        "Age": [28, 32, 25, 35, 30, 27, 40, 29],
        "Department": ["Engineering", "Marketing", "Engineering", "Design", "Marketing", "Engineering", "Design", "Marketing"],
    })
    base_dir = os.path.dirname(os.path.abspath(__file__))
    uploads_dir = os.path.join(base_dir, "data", "uploads")
    os.makedirs(uploads_dir, exist_ok=True)
    path = os.path.join(uploads_dir, "test_sample.csv")
    df.write_csv(path)
    return path


def run_tests():
    print("=" * 60)
    print("ZERO-LLM ENGINE — SMOKE TEST")
    print("=" * 60)

    # ── 1. Upload ───────────────────────────────────────────────
    print("\n[1] Creating + uploading test CSV...")
    csv_path = create_test_csv()
    with open(csv_path, "rb") as f:
        result = handle_upload(f.read(), "test_sample.csv")
    sid = result["session_id"]
    print(f"    Session: {sid}")
    print(f"    Rows: {result['rows']}, Columns: {result['columns']}")

    # ── 2. Intent Parsing Tests ─────────────────────────────────
    test_commands = [
        # (command, expected_operation, expected_column_or_None)
        ("salary ko 10% badhao",       "increase",           "Salary"),
        ("salary ghatao 15%",          "decrease",           "Salary"),
        ("salary 50000 se zyada dikhao","filter",             "Salary"),
        ("name sort karo chota se bada","sort_asc",           "Name"),
        ("total salary batao",         "sum",                "Salary"),
        ("average age nikalo",         "average",            "Age"),
        ("duplicate hatao",            "remove_duplicates",  None),
        ("age column hatao",           "delete_column",      "Age"),
    ]

    print("\n[2] Intent parsing tests:")
    passed = 0
    for cmd, expected_op, expected_col in test_commands:
        intent = parse_intent(sid, cmd)
        op = intent["operation"] if intent else "NONE"
        col = intent.get("column") if intent else None
        op_ok = op == expected_op
        col_ok = (expected_col is None and col is None) or col == expected_col
        status = "✅" if op_ok and col_ok else "❌"
        if op_ok and col_ok:
            passed += 1
        print(f"    {status} '{cmd}' → op={op}, col={col}")

    print(f"    Parsing: {passed}/{len(test_commands)} passed")

    # ── 3. Execution Tests ──────────────────────────────────────
    exec_commands = [
        "salary ko 10% badhao",
        "total salary batao",
        "average salary nikalo",
        "name sort karo chota se bada",
        "sabse bada salary batao",
        "sabse chhota salary batao",
    ]

    print("\n[3] Execution tests:")
    for cmd in exec_commands:
        intent = parse_intent(sid, cmd)
        if intent is None:
            print(f"    ❌ '{cmd}' → unresolved")
            continue
        error = validate_intent(sid, intent)
        if error:
            print(f"    ❌ '{cmd}' → validation: {error}")
            continue
        try:
            result = dispatch(sid, intent)
            print(f"    ✅ '{cmd}' → {result['message']}")
        except Exception as e:
            print(f"    ❌ '{cmd}' → error: {e}")

    # ── 4. Filter test ──────────────────────────────────────────
    print("\n[4] Filter test:")
    intent = parse_intent(sid, "city Mumbai dikhao")
    if intent:
        error = validate_intent(sid, intent)
        if error:
            print(f"    ❌ Validation: {error}")
        else:
            result = dispatch(sid, intent)
            print(f"    ✅ {result['message']}")

    # ── 5. Verify final state ───────────────────────────────────
    print("\n[5] Final data state:")
    pq_path = f"/home/z/my-project/zero-llm-engine/data/sessions/{sid}.parquet"
    if os.path.exists(pq_path):
        df = pl.read_parquet(pq_path)
        print(f"    Rows: {len(df)}, Columns: {df.columns}")
        print(df)

    # ── Cleanup ─────────────────────────────────────────────────
    delete_session(sid)
    if os.path.exists(csv_path):
        os.remove(csv_path)
    print("\n[6] Cleanup done.")

    # ── Server start instructions ───────────────────────────────
    print("\n" + "=" * 60)
    print("To start the server:")
    print("  cd /home/z/my-project/zero-llm-engine")
    print("  source venv/bin/activate")
    print("  uvicorn main:app --host 0.0.0.0 --port 8000 --reload")
    print("=" * 60)


if __name__ == "__main__":
    run_tests()