#!/usr/bin/env python3
"""
Minimal integration test for training endpoints.
"""
import time
import json
import requests

BASE = "http://localhost:8001"

print("1) Start a training job")
resp = requests.post(f"{BASE}/train/start", json={
    "dataset": "./sample_data/train.jsonl",
    "model_id": "unsloth/gemma-3n-E4B-it",
    "prompt_field": "prompt",
    "response_field": "response",
    "epochs": 1,
    "batch_size": 1,
    "gradient_accumulation": 8,
    "use_bf16": True,
    "dry_run": True
})
print(resp.status_code, resp.text)
resp.raise_for_status()
job = resp.json()
job_id = job["job_id"]
print("job_id=", job_id)

print("2) Poll status (10s)")
for _ in range(10):
    s = requests.get(f"{BASE}/train/status/{job_id}")
    print(s.status_code, json.dumps(s.json(), indent=2))
    time.sleep(1)

print("3) Done")