astro-mcp/tests/test_step4_local_astro.py

"""
Step 4: Call the LOCAL astro-mcp server via MCP SSE.
Tests both paths (DB lookup + direct call) and compares angles against ground truth.
"""
from __future__ import annotations

import asyncio
import json
import os
from datetime import timedelta

from mcp import ClientSession
from mcp.client.sse import sse_client

# ── Ground truth from Step 1 (Placidus, Swiss Ephemeris direct) ─────

REFERENCE = {
    "trump": {
        "label": "Donald Trump",
        "birth_datetime": "1946-06-14T10:54:00-04:00",
        "lat": 40.72677,
        "lon": -73.74152,
        # Ground truth: swe.houses(jd, 40.72677, -73.74152, b'P')
        "asc": 150.025101,
        "mc": 54.423678,
    },
    "einstein": {
        "label": "Albert Einstein",
        "birth_datetime": "1879-03-14T11:30:00+00:53",
        "lat": 48.39841,
        "lon": 9.99155,
        "asc": 98.923507,
        "mc": 339.337618,
    },
    "chaka": {
        "label": "Chaka Khan",
        "birth_datetime": "1953-03-23T21:05:00-06:00",
        "lat": 41.87811,
        "lon": -87.62980,
        "asc": 218.923115,
        "mc": 137.470885,
    },
}

TOLERANCE = 0.01

ASTRO_URL = os.environ.get("ASTRO_MCP_URL", "http://127.0.0.1:7016/mcp/sse")


def _payload(result):
    p = getattr(result, "structuredContent", None)
    if isinstance(p, dict):
        return p
    for item in getattr(result, "content", []) or []:
        t = getattr(item, "text", None)
        if isinstance(t, str) and t.strip():
            try:
                d = json.loads(t)
                if isinstance(d, dict):
                    return d
            except Exception:
                continue
    return {}


def _fmt(lon):
    signs = ["Ari", "Tau", "Gem", "Can", "Leo", "Vir",
             "Lib", "Sco", "Sag", "Cap", "Aqu", "Pis"]
    idx = int(lon // 30) % 12
    deg = lon % 30.0
    return f"{signs[idx]} {deg:05.2f} (abs {lon:.6f})"


def check(ref, angles, houses, path_label):
    asc = angles.get("ascendant", {}).get("absolute_lon")
    mc = angles.get("midheaven", {}).get("absolute_lon")
    h1 = houses[0]["absolute_lon"] if houses else None

    asc_diff = abs(asc - ref["asc"]) if asc else float("inf")
    mc_diff = abs(mc - ref["mc"]) if mc else float("inf")
    asc_ok = asc_diff < TOLERANCE
    mc_ok = mc_diff < TOLERANCE

    print(f"\n  [{path_label}]")
    if asc:
        print(f"    ASC: {_fmt(asc)}  exp: {_fmt(ref['asc'])}  diff={asc_diff:.6f}  {'OK' if asc_ok else 'MISMATCH'}")
    else:
        print(f"    ASC: MISSING")
    if mc:
        print(f"    MC:  {_fmt(mc)}  exp: {_fmt(ref['mc'])}  diff={mc_diff:.6f}  {'OK' if mc_ok else 'MISMATCH'}")
    else:
        print(f"    MC:  MISSING")

    if not asc_ok or not mc_ok:
        print(f"    FULL ANGLES: {json.dumps(angles, indent=6, default=str)[:600]}")

    return asc_ok and mc_ok


async def main():
    print("=" * 70)
    print("STEP 4: Local astro-mcp — find where angles get corrupted")
    print(f"  URL: {ASTRO_URL}")
    print("=" * 70)

    url = ASTRO_URL
    if not url.endswith("/mcp/sse"):
        url = url.rstrip("/") + "/mcp/sse"

    all_ok = True

    async with sse_client(url, timeout=30, sse_read_timeout=30) as streams:
        async with ClientSession(*streams, read_timeout_seconds=timedelta(seconds=30)) as session:
            await session.initialize()

            for key, ref in REFERENCE.items():
                print(f"\n{'='*50}")
                print(f"  {ref['label']}")
                print(f"  Input: {ref['birth_datetime']}  lat={ref['lat']}  lon={ref['lon']}")

                # Path A: DB lookup
                chart_a = _payload(await session.call_tool(
                    "calculate_natal_chart_by_id",
                    {"person_id": key, "house_system": "placidus"},
                ))
                if "error" in chart_a:
                    print(f"\n  [DB PATH] ERROR: {chart_a['error']}")
                    all_ok = False
                else:
                    ok = check(ref, chart_a.get("angles", {}), chart_a.get("houses", []), "DB by_id")
                    if not ok:
                        all_ok = False

                # Path B: Direct call
                chart_b = _payload(await session.call_tool(
                    "calculate_natal_chart",
                    {
                        "birth_datetime": ref["birth_datetime"],
                        "latitude": ref["lat"],
                        "longitude": ref["lon"],
                        "house_system": "placidus",
                    },
                ))
                if "error" in chart_b:
                    print(f"\n  [DIRECT PATH] ERROR: {chart_b['error']}")
                    all_ok = False
                else:
                    ok = check(ref, chart_b.get("angles", {}), chart_b.get("houses", []), "DIRECT")
                    if not ok:
                        all_ok = False

    print("\n" + "=" * 70)
    if all_ok:
        print("RESULT: LOCAL ASTRO-MCP IS CORRECT")
    else:
        print("RESULT: LOCAL ASTRO-MCP HAS MISMATCHES")
    print("=" * 70)


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