test_openenv.py

import argparse
import sys

import requests


DEFAULT_BASE_URL = "http://localhost:7860"
TIMEOUT_SECONDS = 10


def print_result(name: str, passed: bool) -> None:
    status = "PASS" if passed else "FAIL"
    print(f"{status} - {name}")


def request_json(method: str, url: str, **kwargs):
    try:
        response = requests.request(method, url, timeout=TIMEOUT_SECONDS, **kwargs)
        print(f"Status Code: {response.status_code}")
        try:
            data = response.json()
            print("Response JSON:", data)
        except ValueError:
            data = None
            print("Response JSON: <invalid JSON>")
            print("Raw Response:", response.text)
        return response, data
    except requests.RequestException as exc:
        print(f"Request error: {exc}")
        return None, None


def has_required_keys(data) -> bool:
    return isinstance(data, dict) and all(key in data for key in ("observation", "reward", "done"))


def test_reset(base_url: str) -> bool:
    print("\nTesting POST /openenv/reset")
    response, data = request_json("POST", f"{base_url}/openenv/reset", json={})

    if response is None:
        print_result("POST /openenv/reset", False)
        return False

    passed = response.status_code == 200 and has_required_keys(data)
    if response.status_code == 405:
        print("Expected 200 but received 405 Method Not Allowed.")
    elif response.status_code != 200:
        print(f"Expected 200 but received {response.status_code}.")
    elif not has_required_keys(data):
        print("Response JSON is missing one or more required keys: observation, reward, done.")

    print_result("POST /openenv/reset", passed)
    return passed


def test_step(base_url: str) -> bool:
    print("\nTesting POST /openenv/step")
    payload = {"action": {"type": "choose_career", "target": "AI Engineer"}}
    response, data = request_json("POST", f"{base_url}/openenv/step", json=payload)

    if response is None:
        print_result("POST /openenv/step", False)
        return False

    passed = response.status_code == 200 and has_required_keys(data)
    if response.status_code != 200:
        print(f"Expected 200 but received {response.status_code}.")
    elif not has_required_keys(data):
        print("Response JSON is missing one or more required keys: observation, reward, done.")

    print_result("POST /openenv/step", passed)
    return passed


def test_state(base_url: str) -> bool:
    print("\nTesting GET /openenv/state")
    response, data = request_json("GET", f"{base_url}/openenv/state")

    if response is None:
        print_result("GET /openenv/state", False)
        return False

    passed = response.status_code == 200 and has_required_keys(data)
    if response.status_code != 200:
        print(f"Expected 200 but received {response.status_code}.")
    elif not has_required_keys(data):
        print("Response JSON is missing one or more required keys: observation, reward, done.")

    print_result("GET /openenv/state", passed)
    return passed


def main() -> int:
    parser = argparse.ArgumentParser(description="Validate OpenEnv FastAPI endpoints.")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL, help=f"Base URL for the API (default: {DEFAULT_BASE_URL})")
    args = parser.parse_args()

    base_url = args.base_url.rstrip("/")
    print(f"Using base URL: {base_url}")

    results = [
        test_reset(base_url),
        test_step(base_url),
        test_state(base_url),
    ]

    overall = all(results)
    print_result("Overall OpenEnv validation", overall)
    return 0 if overall else 1


if __name__ == "__main__":
    sys.exit(main())