Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """Smoke-test SSE progress plus per-hit review artifacts.""" | |
| from __future__ import annotations | |
| import io | |
| import json | |
| import sys | |
| from pathlib import Path | |
| import soundfile as sf | |
| from fastapi.testclient import TestClient | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) | |
| from app import app # noqa: E402 | |
| from synth_generator import generate_test_song # noqa: E402 | |
| def main() -> int: | |
| song = generate_test_song(pattern_name="funk", bars=1, bpm=120, add_bass=False) | |
| buf = io.BytesIO() | |
| sf.write(buf, song.drums_only, song.sr, format="WAV") | |
| buf.seek(0) | |
| client = TestClient(app) | |
| response = client.post( | |
| "/api/jobs", | |
| files={"file": ("funk.wav", buf, "audio/wav")}, | |
| data={"params": json.dumps({"stem": "all", "clustering_mode": "online_preview", "target_min": 2, "target_max": 8})}, | |
| ) | |
| response.raise_for_status() | |
| job_id = response.json()["id"] | |
| final = None | |
| with client.stream("GET", f"/api/jobs/{job_id}/events") as stream: | |
| stream.raise_for_status() | |
| for line in stream.iter_lines(): | |
| if not line or not line.startswith("data: "): | |
| continue | |
| payload = json.loads(line[6:]) | |
| if payload["status"] == "error": | |
| raise RuntimeError(payload.get("error")) | |
| if payload["status"] == "complete": | |
| final = payload | |
| break | |
| assert final is not None, "SSE stream ended without complete event" | |
| hits = final["result"]["hits"] | |
| samples = final["result"]["samples"] | |
| assert hits, "expected review hit rows" | |
| assert samples, "expected representative sample rows" | |
| first_hit_url = hits[0]["url"] | |
| file_response = client.get(first_hit_url) | |
| assert file_response.status_code == 200, first_hit_url | |
| assert file_response.content[:4] == b"RIFF", "review hit should be a WAV file" | |
| print(json.dumps({ | |
| "status": final["status"], | |
| "job_id": job_id, | |
| "hit_count": len(hits), | |
| "sample_count": len(samples), | |
| "first_hit_url": first_hit_url, | |
| }, indent=2)) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |