test_data/enroll.py

"""
Bulk-enroll students into batches using the enroll_requests.json generated
by generate_test_data.py.

Usage:
    python test_data/enroll.py
    python test_data/enroll.py --base-url https://arcshukla-ai-scheduling-platform.hf.space

Requires: requests  (pip install requests  or  uv add requests)
"""

import argparse
import json
import os
import sys

try:
    import requests
except ImportError:
    sys.exit("Install requests first:  pip install requests")

HERE = os.path.dirname(__file__)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default="http://localhost:7861",
                        help="Base URL of the running app (no trailing slash)")
    args = parser.parse_args()
    base = args.base_url.rstrip("/")

    enroll_file = os.path.join(HERE, "enroll_requests.json")
    if not os.path.exists(enroll_file):
        sys.exit(f"Not found: {enroll_file}  – run generate_test_data.py first")

    with open(enroll_file) as f:
        enroll_map: dict[str, list[str]] = json.load(f)

    total_ok = 0
    total_err = 0

    for batch_id, student_ids in enroll_map.items():
        url = f"{base}/api/v1/batches/{batch_id}/enroll"
        payload = {"student_ids": student_ids}
        try:
            r = requests.post(url, json=payload, timeout=30)
            if r.ok:
                print(f"  [OK]  {batch_id}  enrolled {len(student_ids)} students")
                total_ok += 1
            else:
                print(f"  [ERR] {batch_id}  HTTP {r.status_code}: {r.text[:200]}")
                total_err += 1
        except Exception as exc:
            print(f"  [ERR] {batch_id}  {exc}")
            total_err += 1

    print(f"\nDone: {total_ok} batches enrolled, {total_err} errors")


if __name__ == "__main__":
    main()