Baseline.py

# baseline.py — smoke-tests the running CodeRescue API
# Usage: python baseline.py [base_url]
# Default base_url: http://localhost:7860

import sys
import json
import time
import urllib.request
import urllib.error
from env import CodeEnv

BASE_URL = sys.argv[1].rstrip("/") if len(sys.argv) > 1 else "http://localhost:7860"

SAMPLE_CODE = """\
def add(a, b)
    return a + b

print(add(1, 2)
"""

# ── helpers ────────────────────────────────────────────────────────────────────

def get(path: str) -> tuple[int, dict]:
    url = BASE_URL + path
    try:
        with urllib.request.urlopen(url, timeout=10) as r:
            return r.status, json.loads(r.read())
    except urllib.error.HTTPError as e:
        return e.code, {}
    except Exception as e:
        return 0, {"error": str(e)}


def post(path: str, payload: dict) -> tuple[int, dict, float]:
    url  = BASE_URL + path
    data = json.dumps(payload).encode()
    req  = urllib.request.Request(url, data=data,
                                  headers={"Content-Type": "application/json"},
                                  method="POST")
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            elapsed = round(time.time() - t0, 2)
            return r.status, json.loads(r.read()), elapsed
    except urllib.error.HTTPError as e:
        elapsed = round(time.time() - t0, 2)
        return e.code, {}, elapsed
    except Exception as e:
        return 0, {"error": str(e)}, round(time.time() - t0, 2)


def ok(label: str, passed: bool) -> None:
    icon = "✅" if passed else "❌"
    print(f"  {icon}  {label}")

# ── tests ──────────────────────────────────────────────────────────────────────

def test_index():
    print("\n[1] GET /")
    status, body = get("/")
    ok("status 200",          status == 200)
    ok("name == CodeRescue API", body.get("name") == "CodeRescue API")
    ok("status field == running", body.get("status") == "running")


def test_health():
    print("\n[2] GET /health")
    status, body = get("/health")
    ok("status 200",       status == 200)
    ok("status == ok",     body.get("status") == "ok")


def test_analyze_no_code():
    print("\n[3] POST /analyze — empty payload")
    status, body, elapsed = post("/analyze", {})
    ok("status 400",       status == 400)
    ok("error field present", "error" in body)


def test_analyze():
    print("\n[4] POST /analyze — real code")
    status, body, elapsed = post("/analyze", {"code": SAMPLE_CODE})
    ok(f"status 200 (got {status})",   status == 200)
    ok("explanation present",          bool(body.get("explanation")))
    ok("fixed_code present",           bool(body.get("fixed_code")))
    ok("language present",             bool(body.get("language")))
    ok("confidence present",           body.get("confidence") in ("high", "medium", "low"))
    print(f"       Response Time: {elapsed}s")
    if body.get("time_taken") is not None:
        print(f"       Server time_taken: {body['time_taken']}s")


def test_reset():
    print("\n[5] POST /reset")
    status, body, _ = post("/reset", {})
    ok("status 200",           status == 200)
    ok("status == success",    body.get("status") == "success")


def test_openenv_reset():
    print("\n[6] POST /openenv/reset")
    status, body, _ = post("/openenv/reset", {})
    ok("status 200",           status == 200)
    ok("status == success",    body.get("status") == "success")


def test_env():
    print("\n[7] OpenEnv test")
    env = CodeEnv()

    state = env.reset()
    print("State:", state)
   # simulate agent decision (baseline policy)
    action = {"fixed_code": state["code"].replace("/0", "/1")}
    state, reward, done, _ = env.step(action)

    print("Reward:", reward)
    ok("reward is a float",    isinstance(reward, float))
    ok("done is True",         done is True)
    ok("state has code key",   "code" in state)


# ── main ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print(f"🔍 Baseline — target: {BASE_URL}")
    test_index()
    test_health()
    test_analyze_no_code()
    test_analyze()
    test_reset()
    test_openenv_reset()
    test_env()
    print("\nDone.\n")