scripts/test_progress_contract.py

#!/usr/bin/env python3
"""Smoke test the backend job progress contract."""
from __future__ import annotations

import sys
import tempfile
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import numpy as np
import soundfile as sf
from fastapi.testclient import TestClient

from app import app


def make_fixture(path: Path, sr: int = 44100) -> None:
    duration = 1.4
    y = np.zeros(int(sr * duration), dtype=np.float32)
    for onset in [0.1, 0.35, 0.7, 1.05]:
        start = int(onset * sr)
        n = int(0.045 * sr)
        tone = np.sin(2 * np.pi * 180 * np.arange(n) / sr).astype(np.float32)
        env = np.exp(-np.linspace(0, 6, n)).astype(np.float32)
        y[start:start+n] += tone * env * 0.8
    sf.write(path, y, sr)


def main() -> None:
    client = TestClient(app)
    with tempfile.TemporaryDirectory() as tmp:
        src = Path(tmp) / "progress.wav"
        make_fixture(src)
        with src.open("rb") as handle:
            response = client.post(
                "/api/jobs",
                files={"file": (src.name, handle, "audio/wav")},
                data={"params": '{"stem":"all","clustering_mode":"online_preview","target_min":2,"target_max":6}'},
            )
        assert response.status_code == 202, response.text
        job = response.json()
        assert "progress" in job, job
        assert job["progress"]["fraction"] >= 0
        job_id = job["id"]
        for _ in range(900):
            time.sleep(0.1)
            job = client.get(f"/api/jobs/{job_id}").json()
            progress = job.get("progress") or {}
            assert "fraction" in progress, progress
            assert 0 <= progress["fraction"] <= 1, progress
            assert progress.get("basis"), progress
            if job["status"] == "complete":
                assert progress["fraction"] == 1.0, progress
                break
        else:
            raise AssertionError("job did not complete")
        stages = job.get("stages") or []
        assert all("progress" in stage for stage in stages), stages
    print({"status": "ok", "progress": job["progress"], "stage_count": len(job.get("stages") or [])})


if __name__ == "__main__":
    main()