test_v2_tools.py

"""
v2.0 comprehensive live test — all 32 new tools against real SPSS 27.

Run: python test_v2_tools.py
"""
import asyncio
import os
import os.path

from spss_mcp.server import (
    # Group A — data management
    spss_recode, spss_autorecode, spss_compute, spss_if_transform,
    spss_select_if, spss_sort_cases, spss_weight, spss_filter,
    spss_missing_values, spss_variable_labels, spss_value_labels,
    spss_formats, spss_declare_variables, spss_save_file,
    # Group B — restructuring
    spss_match_files, spss_add_files, spss_star_join, spss_aggregate,
    spss_rank, spss_flip, spss_cases_to_vars, spss_vars_to_cases,
    # Group C — reporting
    spss_export_output, spss_graph_scatter, spss_graph_histogram,
    spss_graph_boxplot, spss_graph_bar, spss_custom_tables,
    # Group D — advanced
    spss_quick_cluster, spss_proximities, spss_arima, spss_ratio_statistics,
)

T = os.path.abspath("test_data.sav")     # 5 rows: age, income, gender
B = os.path.abspath("big_data.sav")       # 30 rows: age, income, gender, group
OUT = lambda name: os.path.abspath(f"v2test_{name}.sav")


def ensure_test_data() -> None:
    """Create big_data.sav if missing (30 synthetic rows: age, income, gender, group)."""
    if os.path.exists(B):
        return
    import pandas
    import pyreadstat

    ages = [23, 25, 27, 28, 30, 31, 33, 35, 36, 38,
            40, 41, 43, 45, 47, 48, 50, 52, 54, 55,
            57, 58, 60, 62, 63, 65, 67, 68, 70, 72]
    incomes = [2200, 2800, 3100, 2500, 4200, 3900, 4600, 5000, 4400, 5500,
               6100, 5800, 6700, 7300, 6900, 7800, 8200, 7500, 8900, 9100,
               8600, 9800, 10400, 9500, 11200, 10800, 12500, 11900, 13200, 14000]
    genders = [1, 2] * 15
    groups = [1, 1, 1, 2, 2, 2, 1, 1, 2, 2,
              1, 2, 1, 2, 1, 2, 1, 1, 2, 2,
              1, 2, 1, 1, 2, 1, 2, 2, 1, 2]
    df = pandas.DataFrame({
        "age": ages, "income": incomes, "gender": genders, "group": groups,
    })
    df["gender"] = df["gender"].astype("float64")
    df["group"] = df["group"].astype("float64")
    pyreadstat.write_sav(df, B)


def is_pass(result: str) -> bool:
    # NOTE: "Error" appears in legit SPSS output (ANOVA "Error" row, ARIMA
    # "Errors for Autocorrelations") — only flag failures, not table content.
    return not result.startswith("Error") and "was not created" not in result


async def main():
    ensure_test_data()
    results = []

    async def check(name, coro):
        try:
            r = await coro
            ok = is_pass(r)
        except Exception as e:
            r, ok = str(e), False
        results.append((name, ok))
        print(f"{'PASS' if ok else 'FAIL'} {name}")

    # ─── Group A ────────────────────────────────────────────────────────────
    await check("recode", spss_recode(
        file_path=B, variables=["income"],
        rules=[{"old": [0, 3000], "new": 1}, {"old": "ELSE", "new": 2}],
        output_path=OUT("recode"), ctx=None))
    await check("autorecode", spss_autorecode(
        file_path=B, variables=["gender"], new_names=["gender_n"],
        output_path=OUT("autorecode"), ctx=None))
    await check("compute", spss_compute(
        file_path=B, target_variable="score", expression="income/1000 + age",
        output_path=OUT("compute"), ctx=None))
    await check("if_transform", spss_if_transform(
        file_path=B, condition="age >= 30", target_variable="agegrp",
        expression="2", output_path=OUT("if"), ctx=None))
    await check("select_if", spss_select_if(
        file_path=B, condition="age >= 25", output_path=OUT("select"), ctx=None))
    await check("sort_cases", spss_sort_cases(
        file_path=B, sort_keys=[{"variable": "age", "order": "D"}],
        output_path=OUT("sort"), ctx=None))
    await check("weight", spss_weight(file_path=B, weight_variable="income", ctx=None))
    await check("filter", spss_filter(file_path=B, filter_variable="gender", ctx=None))
    await check("missing_values", spss_missing_values(
        file_path=B, missing_spec={"age": "99"}, output_path=OUT("miss"), ctx=None))
    await check("variable_labels", spss_variable_labels(
        file_path=B, labels={"age": "Respondent age in years"},
        output_path=OUT("varlab"), ctx=None))
    await check("value_labels", spss_value_labels(
        file_path=B, value_labels={"gender": {1: "Male", 2: "Female"}},
        output_path=OUT("vallab"), ctx=None))
    await check("formats", spss_formats(
        file_path=B, formats={"income": "F8.1"}, output_path=OUT("fmt"), ctx=None))
    await check("declare_variables", spss_declare_variables(
        file_path=B, numeric_vars={"newv": "F8.2"},
        output_path=OUT("declare"), ctx=None))
    await check("save_file", spss_save_file(
        file_path=B, output_path=OUT("save"), keep=["age", "income"], ctx=None))

    # ─── Group B ────────────────────────────────────────────────────────────
    # second file for merge tests
    from spss_mcp.spss_runner import run_syntax
    await run_syntax(
        f"GET FILE='{B}'.\nSAVE OUTFILE='{os.path.abspath('v2test_second.sav')}'.",
        save_viewer_output=False)

    second = os.path.abspath("v2test_second.sav")
    await check("match_files", spss_match_files(
        file_paths=[B, second], key_variables=["age"],
        output_path=OUT("match"), ctx=None))
    await check("add_files", spss_add_files(
        file_paths=[B, second], output_path=OUT("add"), ctx=None))
    await check("star_join", spss_star_join(
        base_file=B, join_file=second, base_key="age", join_key="age",
        output_path=OUT("join"), ctx=None))
    await check("aggregate", spss_aggregate(
        file_path=B, break_variables=["gender"],
        aggregations=[{"new_var": "inc_mean", "function": "MEAN", "source": "income"}],
        output_path=OUT("agg"), ctx=None))
    await check("rank", spss_rank(
        file_path=B, variables=["income"], order="D", rank_into="Rinc",
        output_path=OUT("rank"), ctx=None))
    await check("flip", spss_flip(file_path=T, output_path=OUT("flip"), ctx=None))
    await check("vars_to_cases", spss_vars_to_cases(
        file_path=T, make_specs=[{"new_var": "value", "variables": ["age", "income"]}],
        index_name="variable", id_name="case_id", output_path=OUT("long"), ctx=None))
    await check("cases_to_vars", spss_cases_to_vars(
        file_path=OUT("long"), id_variables=["case_id"],
        index_variables=["variable"], output_path=OUT("wide"), ctx=None))

    # ─── Group C ────────────────────────────────────────────────────────────
    await check("graph_scatter", spss_graph_scatter(
        file_path=B, x_variable="age", y_variable="income", ctx=None))
    await check("graph_histogram", spss_graph_histogram(
        file_path=B, variable="income", normal_curve=True, ctx=None))
    await check("graph_boxplot", spss_graph_boxplot(
        file_path=B, variable="income", categorical_variable="gender", ctx=None))
    await check("graph_bar", spss_graph_bar(
        file_path=B, variable="income", statistic="MEAN", by_variable="gender", ctx=None))
    await check("custom_tables", spss_custom_tables(
        file_path=B, rows=["income"], columns=["gender"], statistics=["MEAN"], ctx=None))
    await check("export_output", spss_export_output(
        file_path=B, procedures_syntax="FREQUENCIES VARIABLES=gender.",
        output_path=os.path.abspath("v2test_export.html"),
        export_format="HTML", ctx=None))

    # ─── Group D ────────────────────────────────────────────────────────────
    await check("quick_cluster", spss_quick_cluster(
        file_path=B, variables=["age", "income"], n_clusters=3,
        save_cluster_variable="cl", print_anova=True, ctx=None))
    await check("proximities", spss_proximities(
        file_path=B, variables=["age", "income"], measure="EUCLID", ctx=None))
    await check("arima", spss_arima(
        file_path=B, dependent_variable="income", p=1, d=0, q=0, ctx=None))
    await check("ratio_statistics", spss_ratio_statistics(
        file_path=B, numerator="income", denominator="age",
        group_variable="gender", ctx=None))

    # ─── Summary ─────────────────────────────────────────────────────────────
    passed = sum(1 for _, ok in results if ok)
    print(f"\n{'=' * 50}")
    print(f"v2.0 LIVE TEST: {passed}/{len(results)} PASS ({100 * passed / len(results):.0f}%)")
    failed = [n for n, ok in results if not ok]
    if failed:
        print("Failed:", ", ".join(failed))
    print(f"Total tools in server: 69 (37 original + 32 new)")


if __name__ == "__main__":
    asyncio.run(main())